RGB depth

app.py

@@ -116,7 +116,7 @@ with block:
     def on_depth_reconstruction_btn_click(inputs):
         print("depth reconstruction")
         image = inputs[input_image]
-        path = dpt.generate_obj_masks(image=image, n_samples=inputs[n_samples], cube_size=inputs[cube_size], masks=inputs[masks])
+        path = dpt.generate_obj_rgb(image=image, n_samples=inputs[n_samples], cube_size=inputs[cube_size]) #
         return {pcl_figure: path}
     depth_reconstruction_btn.click(on_depth_reconstruction_btn_click, components, [pcl_figure], queue=False)
 

inference.py

@@ -36,7 +36,7 @@ class DepthPredictor:
         output = prediction.cpu().numpy()
         formatted = (output * 255 / np.max(output)).astype('uint8')
         #img = Image.fromarray(formatted)
-        return output
+        return formatted
     
     def generate_pcl(self, image):
         print(np.array(image).shape)