3d segmentation reconstruction

app.py

@@ -9,18 +9,26 @@ import numpy as np
 def snap(image, video):
     depth_predictor = DepthPredictor()
     depth_result = depth_predictor.predict(image)
-    gltf_path = create_3d_obj(np.array(image), depth_result)
-    #segment_predictor = SegmentPredictor()
-    #sam_result = segment_predictor.predict(image)
-    return [image, gltf_path, gltf_path]#[depth_result, gltf_path, gltf_path]
+    rgb_gltf_path = create_3d_obj(np.array(image), depth_result, path='./rgb.gltf')
+    
+    segment_predictor = SegmentPredictor()
+    sam_result = segment_predictor.predict(image)
+    seg_gltf_path = create_3d_obj(np.array(sam_result), depth_result, path='./seg.gltf')
+
+    return [image, depth_result, sam_result, rgb_gltf_path, seg_gltf_path]#[depth_result, gltf_path, gltf_path]
 
 
 demo = gr.Interface(
     snap,
     inputs=[gr.Image(source="webcam", tool=None, type="pil"),
             gr.Video(source="webcam")],
-    outputs=[gr.Image(label="predicted depth", type="pil"),
-             gr.Model3D(label="3d mesh reconstruction", clear_color=[1.0, 1.0, 1.0, 1.0]), gr.File(label="3d gLTF")]
+    outputs=[gr.Image(label="RGB", type="pil"),
+             gr.Image(label="predicted depth", type="pil"),
+             gr.Image(label="predicted segmentation", type="pil"),
+             gr.Model3D(label="3D mesh reconstruction - RGB",
+                        clear_color=[1.0, 1.0, 1.0, 1.0]),
+             gr.Model3D(label="3D mesh reconstruction - Segmented",
+                        clear_color=[1.0, 1.0, 1.0, 1.0])]
 )
 
 if __name__ == "__main__":

utils.py

@@ -2,7 +2,7 @@ import numpy as np
 import open3d as o3d
 
 
-def create_3d_obj(rgb_image, depth_image, depth=10):
+def create_3d_obj(rgb_image, depth_image, depth=10, path='./image.gltf'):
     depth_o3d = o3d.geometry.Image(depth_image)
     image_o3d = o3d.geometry.Image(rgb_image)
     rgbd_image = o3d.geometry.RGBDImage.create_from_color_and_depth(
@@ -47,7 +47,7 @@ def create_3d_obj(rgb_image, depth_image, depth=10):
     # mesh.remove_vertices_by_mask(vertices_to_remove)
     bbox = pcd.get_axis_aligned_bounding_box()
     mesh_crop = mesh.crop(bbox)
-    gltf_path = './image.gltf'
+    gltf_path = path
     o3d.io.write_triangle_mesh(
         gltf_path, mesh_crop, write_triangle_uvs=True)
     return gltf_path