UI fix

app.py

@@ -24,46 +24,36 @@ def produce_3d_reconstruction(image):
 def produce_point_cloud(depth_map, segmentation_map):
     return point_cloud(np.array(segmentation_map), depth_map)
 
-def snap(image, depth_map, segmentation_map, point_cloud, video):
-    if depth_map:
-        depth_result = produce_depth_map(image)
-    else:
-        depth_result = None
-        
-    if segmentation_map:
-        sam_result = produce_segmentation_map(image)
-    else:
-        sam_result = None
-        
-    if point_cloud:
-        point_cloud_fig = produce_point_cloud(depth_result, sam_result)
-    else:
-        point_cloud_fig = None
-        
+def snap(image, depth_map, segmentation_map, video):
+    depth_result = produce_depth_map(image) if depth_map else None
+    sam_result = produce_segmentation_map(image) if segmentation_map else None
+    rgb_gltf_path = produce_3d_reconstruction(image) if depth_map else None
+    point_cloud_fig = produce_point_cloud(depth_result, sam_result) if segmentation_map else None
+    
     if video:
         # Add video processing here if needed
         pass
 
-    return [image, depth_result, sam_result, point_cloud_fig]
+    return [image, depth_result, sam_result, rgb_gltf_path, point_cloud_fig]
 
 # Interface inputs
 image_input = gr.Image(source="webcam", tool=None, label="Input Image", type="pil")
 depth_map_button = gr.Button(label="Produce Depth Map", value=False)
 segmentation_map_button = gr.Button(label="Produce Segmentation Map", value=False)
-point_cloud_button = gr.Button(label="Produce Point Cloud", value=False)
 video_input = gr.Video(source="webcam")
 
 # Interface outputs
 output_image = gr.Image(label="RGB")
 output_depth_map = gr.Image(label="Predicted Depth")
 output_segmentation_map = gr.Image(label="Predicted Segmentation")
+output_3d_reconstruction = gr.Model3D(label="3D mesh reconstruction - RGB", clear_color=[1.0, 1.0, 1.0, 1.0])
 output_point_cloud = gr.Plot(label="Point Cloud")
 
 # Interface
 demo = gr.Interface(
     snap,
-    inputs=[image_input, depth_map_button, segmentation_map_button, point_cloud_button, video_input],
-    outputs=[output_image, output_depth_map, output_segmentation_map, None, output_point_cloud]
+    inputs=[image_input, depth_map_button, segmentation_map_button, video_input],
+    outputs=[output_image, output_depth_map, output_segmentation_map, output_3d_reconstruction, output_point_cloud]
 )
 
 if __name__ == "__main__":