instance on cubes

app.py

@@ -34,7 +34,8 @@ with block:
         with gr.Row():
             input_image = gr.Image(label='Input', height=512, type='pil')
             masks_annotated_image = gr.AnnotatedImage(label='Segments', height=512)
-            pcl_figure = gr.Plot(label='3D Reconstruction')
+            pcl_figure = gr.Model3D(label="3-D Reconstruction", clear_color=[1.0, 1.0, 1.0, 1.0])
+            #gr.Plot(label='3D Reconstruction')
             #cutout_galary = gr.Gallery(label='Cutouts', object_fit='contain', height=512)
         with gr.Row():
             with gr.Column(scale=1):
@@ -53,8 +54,8 @@ with block:
         print("depth reconstruction")
         image = inputs[raw_image]
         # depth reconstruction
-        fig = dpt.generate_fig2(image)
-        return {pcl_figure: fig}
+        path = dpt.generate_obj(image)
+        return {pcl_figure: path}
     
     depth_reconstruction_btn.click(on_depth_reconstruction_btn_click, components, [pcl_figure], queue=False)
     

inference.py

@@ -69,6 +69,26 @@ class DepthPredictor:
         ax = fig.add_subplot(111, projection='3d')
         ax.scatter(points,size=0.01, c=colors, marker='o')
         return fig
+    
+    def generate_obj(self, image):
+        # Step 1: Create a point cloud
+        point_cloud, color_array = self.generate_pcl(image)
+        # sample 1000 points
+        idxs = np.random.choice(len(point_cloud), 1000)
+        point_cloud = point_cloud[idxs]
+        color_array = color_array[idxs]
+        # Create a mesh to hold the colored cubes
+        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.01, height=0.01, depth=0.01)
+            cube.translate(point)
+            cube.paint_uniform_color(color)
+            mesh += cube
+        # Save the mesh to an .obj file
+        output_file = "./cloud.obj"
+        o3d.io.write_triangle_mesh(output_file, mesh)
+        return output_file