masked pointcloud

app.py

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

inference.py

@@ -10,6 +10,7 @@ import open3d as o3d
 import pandas as pd
 import plotly.express as px
 import matplotlib.pyplot as plt
+from utils import PCL
 
 class DepthPredictor:
     def __init__(self):
@@ -103,10 +104,6 @@ class DepthPredictor:
         # 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):
-            #if len(mask) == len(point_cloud):
-            #    mask = mask.ravel()
-            #else:
-            #    mask = mask.ravel()[:-1]
             mask = mask.ravel()
             point_cloud_subset, color_array_subset = point_cloud[mask], color_array[mask]
             idxs = np.random.choice(len(point_cloud_subset), int(n_samples))
@@ -121,6 +118,25 @@ class DepthPredictor:
         o3d.io.write_triangle_mesh(output_file, mesh)
         return output_file
     
+    def generate_obj_masks2(self, image, masks, cube_size):
+        # Generate a point cloud
+        depth = self.predict(image)
+        image = np.array(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):
+            points, _ = PCL(mask, depth)
+            for point in points:
+                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
+    
 
 
 

utils.py

@@ -193,4 +193,24 @@ def PCL3(image):
     # Step 6: Create a 3D scatter plot using Plotly Express
     #fig = px.scatter_3d(df, x='x', y='y', z='z', color=colors, size="sizes")
     
-    return fig
+    return fig
+
+
+def PCL(mask, depth):
+    assert mask.shape == depth.shape
+    assert type(mask) == np.ndarray
+    assert type(depth) == np.ndarray
+    rgb_mask = np.zeros((mask.shape[0], mask.shape[1], 3))
+    rgb_mask[mask] = (1.0, 0.0, 0.0)
+    depth_o3d = o3d.geometry.Image(depth)
+    image_o3d = o3d.geometry.Image(mask)
+    rgbd_image = o3d.geometry.RGBDImage.create_from_color_and_depth(image_o3d, depth_o3d, convert_rgb_to_intensity=False)
+    # Step 3: Create a PointCloud from the RGBD image
+    pcd = o3d.geometry.PointCloud.create_from_rgbd_image(rgbd_image, o3d.camera.PinholeCameraIntrinsic(o3d.camera.PinholeCameraIntrinsicParameters.PrimeSenseDefault))
+    # Step 4: Convert PointCloud data to a NumPy array
+    points = np.asarray(pcd.points)
+    colors = np.asarray(pcd.colors)
+    mask = (colors[:, 0] == 1.0)
+    points = points[mask]
+    colors = colors[mask]
+    return points, colors