cube size slider, masked pointcloud

app.py

@@ -35,7 +35,9 @@ with block:
             input_image = gr.Image(label='Input', type='pil', source='webcam', tool=None)
             masks_annotated_image = gr.AnnotatedImage(label='Segments', height=512)
             pcl_figure = gr.Model3D(label="3-D Reconstruction", clear_color=[1.0, 1.0, 1.0, 1.0])
-            n_samples = gr.Slider(minimum=1e3, maximum=1e6, step=1e3, default=1e3, label='Number of Samples')
+            with gr.Column():
+                n_samples = gr.Slider(minimum=1e3, maximum=1e6, step=1e3, default=1e3, label='Number of Samples')
+                cube_size = gr.Slider(minimum=0.000001, maximum=0.001, step=0.000001, default=0.00001, label='Cube size')
         with gr.Row():
             with gr.Column(scale=1):
                 with gr.Row():
@@ -48,7 +50,7 @@ with block:
     # components
     components = {point_coords, point_labels, raw_image, masks, cutout_idx, input_image,
                   point_label_radio, text, reset_btn, sam_sgmt_everything_btn,
-                  sam_decode_btn, depth_reconstruction_btn, masks_annotated_image, n_samples}
+                  sam_decode_btn, depth_reconstruction_btn, masks_annotated_image, n_samples, cube_size}
     
     # event - init coords
     def on_reset_btn_click(raw_image):
@@ -86,7 +88,7 @@ with block:
     def on_depth_reconstruction_btn_click(inputs):
         print("depth reconstruction")
         image = inputs[input_image]
-        path = dpt.generate_obj(image, inputs[n_samples])
+        path = dpt.generate_obj_masks(image, inputs[n_samples], inputs[cube_size])
         return {pcl_figure: path}
     depth_reconstruction_btn.click(on_depth_reconstruction_btn_click, components, [pcl_figure], queue=False)
 

inference.py

@@ -70,7 +70,7 @@ class DepthPredictor:
         ax.scatter(points,size=0.01, c=colors, marker='o')
         return fig
     
-    def generate_obj(self, image, n_samples):
+    def generate_obj_rgb(self, image, n_samples, cube_size):
         # Step 1: Create a point cloud
         point_cloud, color_array = self.generate_pcl(image)
         #point_cloud, color_array = point_cloud[mask.ravel()[:-1]], color_array[mask.ravel()[:-1]]
@@ -82,7 +82,7 @@ class DepthPredictor:
         mesh = o3d.geometry.TriangleMesh()
         # Create cubes and add them to the mesh
         for point, color in zip(point_cloud, color_array):
-            cube = o3d.geometry.TriangleMesh.create_box(width=0.000001, height=0.000001, depth=0.000001)
+            cube = o3d.geometry.TriangleMesh.create_box(width=cube_size, height=cube_size, depth=cube_size)
             cube.translate(-point)
             cube.paint_uniform_color(color)
             mesh += cube
@@ -91,6 +91,26 @@ class DepthPredictor:
         o3d.io.write_triangle_mesh(output_file, mesh)
         return output_file
 
+    def generate_obj_masks(self, image, n_samples, masks, cube_size):
+        # Generate a point cloud
+        point_cloud, color_array = self.generate_pcl(image)
+        mesh = o3d.geometry.TriangleMesh()
+        # Create cubes and add them to the mesh
+        cs = [(255,0,0),(0,255,0),(0,0,255)]
+        for c,(mask, _) in zip(cs, masks):
+            mask = mask.ravel()[:-1]
+            point_cloud_subset, color_array_subset = point_cloud[mask], color_array[mask]
+            idxs = np.random.choice(len(point_cloud_subset), int(n_samples))
+            point_cloud_subset = point_cloud_subset[idxs]
+            for point in point_cloud_subset:
+                cube = o3d.geometry.TriangleMesh.create_box(width=cube_size, height=cube_size, depth=cube_size)
+                cube.translate(-point)
+                cube.paint_uniform_color(c)
+                mesh += cube
+        # Save the mesh to an .obj file
+        output_file = "./cloud.obj"
+        o3d.io.write_triangle_mesh(output_file, mesh)
+        return output_file