depth ui

app.py

@@ -5,10 +5,11 @@ import cv2
 from PIL import Image
 import torch
 from inference import SegmentPredictor
+from utils import generate_PCL
 
 
 
-sam = SegmentPredictor() #service.get_sam(configs.model_type, configs.model_ckpt_path, configs.device)
+sam = SegmentPredictor()
 red = (255,0,0)
 blue = (0,0,255)
 annos = []
@@ -32,20 +33,30 @@ with block:
         with gr.Row():
             input_image = gr.Image(label='Input', height=512, type='pil')
             masks_annotated_image = gr.AnnotatedImage(label='Segments', height=512)
-            cutout_galary = gr.Gallery(label='Cutouts', object_fit='contain', height=512)
+            pcl_figure = gr.Plot(label='3D Reconstruction')
+            #cutout_galary = gr.Gallery(label='Cutouts', object_fit='contain', height=512)
         with gr.Row():
             with gr.Column(scale=1):
                 with gr.Row():
                     point_label_radio = gr.Radio(label='Point Label', choices=[1,0], value=1)
                     text = gr.Textbox(label='Mask Name')
+                    reset_btn = gr.Button('New Mask')
                 sam_sgmt_everything_btn = gr.Button('Segment Everything!', variant = 'primary')
                 sam_decode_btn = gr.Button('Predict using points!', variant = 'primary')
-                reset_btn = gr.Button('Reset')
-
+                depth_reconstruction_btn = gr.Button('Depth Reconstruction', variant = 'primary')
     # components
     components = {point_coords, point_labels, raw_image, masks, cutout_idx, input_image,
                   point_label_radio, text, reset_btn, sam_sgmt_everything_btn,
-                  sam_decode_btn, masks_annotated_image}
+                  sam_decode_btn, depth_reconstruction_btn, masks_annotated_image}
+    def on_depth_reconstruction_btn_click(inputs):
+        print("depth reconstruction")
+        image = inputs[raw_image]
+        # depth reconstruction
+        fig = generate_PCL(image)
+        return {pcl_figure: fig}
+    
+    depth_reconstruction_btn.click(on_depth_reconstruction_btn_click, components, [pcl_figure], queue=False)
+    
     # event - init coords
     def on_reset_btn_click(raw_image):
         return raw_image, point_coords_empty(), point_labels_empty(), None, []
@@ -79,7 +90,6 @@ with block:
         inputs[masks].append((generated_mask, inputs[text]))
         return {masks_annotated_image: (image, inputs[masks])}
     sam_decode_btn.click(on_click_sam_dencode_btn, components, [masks_annotated_image, masks, cutout_idx], queue=True)
-    #sam_sgmt_everything_btn.click(on_sam_sgmt_everything_click, components, [masks_annotated_image, masks, cutout_idx], queue=True)
 
 
 if __name__ == '__main__':

app_legacy.py

@@ -2,7 +2,7 @@ import gradio as gr
 from segment_anything import SamAutomaticMaskGenerator, sam_model_registry
 import supervision as sv
 from inference import DepthPredictor, SegmentPredictor
-from utils import create_3d_obj, create_3d_pc, point_cloud
+from utils import create_3d_obj, create_3d_pc, point_cloud, generate_PCL
 import numpy as np
 
 def produce_depth_map(image):

utils.py

@@ -4,7 +4,7 @@ import open3d as o3d
 import plotly.express as px
 import numpy as np
 import pandas as pd
-
+from inference import DepthPredictor
 
 def create_3d_obj(rgb_image, depth_image, depth=10, path='./image.gltf'):
     depth_o3d = o3d.geometry.Image(depth_image)
@@ -139,4 +139,17 @@ def array_PCL(rgb_image, depth_image):
     xx_rgb = ((rgb_image[:, 0] * FX_RGB) / rgb_image[:, 2] + CX_RGB + width / 2).astype(int).clip(0, width - 1)
     yy_rgb = ((rgb_image[:, 1] * FY_RGB) / rgb_image[:, 2] + CY_RGB).astype(int).clip(0, height - 1)
     colors = rgb_image[yy_rgb, xx_rgb]/255
-    return pcd, colors
+    return pcd, colors
+
+def generate_PCL(image):
+    depth_predictor = DepthPredictor()
+    depth_result = depth_predictor.predict(image)
+    pcd, colors = array_PCL(image, depth_result)
+    fig = px.scatter_3d(x=pcd[:, 0], y=pcd[:, 1], z=pcd[:, 2], color=colors, size_max=0.1)
+    return fig
+
+
+def plot_PCL(rgb_image, depth_image):
+    pcd, colors = array_PCL(rgb_image, depth_image)
+    fig = px.scatter_3d(x=pcd[:, 0], y=pcd[:, 1], z=pcd[:, 2], color=colors, size_max=0.1)
+    return fig