default value and mask reconstruction

app.py

@@ -74,16 +74,16 @@ with block:
                 )
                 with gr.Row():
                     max_depth = gr.Slider(
-                        minimum=0, maximum=10, step=0.01, default=1, label="Max Depth"
+                        minimum=0, maximum=10, value=3, step=0.01, label="Max Depth"
                     )
                     min_depth = gr.Slider(
-                        minimum=0, maximum=10, step=0.01, default=0.1, label="Min Depth"
+                        minimum=0, maximum=10, step=0.01, value=1, label="Min Depth"
                     )
                     n_samples = gr.Slider(
                         minimum=1e3,
                         maximum=1e6,
                         step=1e3,
-                        default=1e3,
+                        value=1e5,
                         label="Number of Samples",
                     )
                     cube_size = gr.Slider(
@@ -94,7 +94,10 @@ with block:
                         label="Cube size",
                     )
                     depth_reconstruction_btn = gr.Button(
-                        "Depth Reconstruction", variant="primary"
+                        "3D Reconstruction", variant="primary"
+                    )
+                    depth_reconstruction_mask_btn = gr.Button(
+                        "Mask Reconstruction", variant="primary"
                     )
 
                 sam_decode_btn = gr.Button("Predict using points!", variant="primary")
@@ -269,6 +272,22 @@ with block:
         on_depth_reconstruction_btn_click, components, [pcl_figure], queue=False
     )
 
+    def on_depth_reconstruction_mask_btn_click(inputs):
+        print("depth reconstruction")
+        path = dpt.generate_obj_masks(
+            image=inputs[input_image],
+            cube_size=inputs[cube_size],
+            n_samples=inputs[n_samples],
+            masks=inputs[masks],
+            min_depth=inputs[min_depth],
+            max_depth=inputs[max_depth],
+        )
+        return {pcl_figure: path}
+
+    depth_reconstruction_mask_btn.click(
+        on_depth_reconstruction_mask_btn_click, components, [pcl_figure], queue=False
+    )
+
     def on_sam_sgmt_everything_btn_click(inputs):
         print("segmenting everything")
         image = inputs[input_image]

inference.py

@@ -229,16 +229,16 @@ class DepthPredictor:
     ):
         # Generate a point cloud
         depth = self.predict(image)
-        # depth = map_image_range(depth, min_depth, max_depth)
+        depth = map_image_range(depth, min_depth, max_depth)
         image = np.array(image)
         mesh = o3d.geometry.TriangleMesh()
         # Create cubes and add them to the mesh
         print(len(masks))
-        cs = [(255, 0, 0), (0, 255, 0), (0, 0, 255)]
+        cs = [(1, 0, 0), (0, 1, 0), (0, 0, 1)]
         for c, (mask, _) in zip(cs, masks):
             points, _ = PCL(mask, depth)
-            # idxs = np.random.choice(len(points), int(n_samples))
-            # points = points[idxs]
+            idxs = np.random.choice(len(points), int(n_samples))
+            points = points[idxs]
             for point in points:
                 cube = o3d.geometry.TriangleMesh.create_box(
                     width=cube_size, height=cube_size, depth=cube_size