masking

app.py

@@ -55,7 +55,9 @@ with block:
         print("depth reconstruction")
         image = inputs[raw_image]
         # depth reconstruction
-        path = dpt.generate_obj(image, inputs[n_samples])
+        mask = inputs[masks][0][0]
+        path = dpt.generate_obj(image, inputs[n_samples], mask)
+
         return {pcl_figure: path}
     
     depth_reconstruction_btn.click(on_depth_reconstruction_btn_click, components, [pcl_figure], queue=False)

inference.py

@@ -70,9 +70,10 @@ class DepthPredictor:
         ax.scatter(points,size=0.01, c=colors, marker='o')
         return fig
     
-    def generate_obj(self, image, n_samples):
+    def generate_obj(self, image, n_samples, mask):
         # Step 1: Create a point cloud
         point_cloud, color_array = self.generate_pcl(image)
+        point_cloud, color_array = point_cloud[mask.ravel()], color_array[mask.ravel()]
         # sample 1000 points
         idxs = np.random.choice(len(point_cloud), int(n_samples))
         point_cloud = point_cloud[idxs]