UI first try

app.py

@@ -1,48 +1,88 @@
+import os
 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
 import numpy as np
+import cv2 
+from PIL import Image
+import torch
+from inference import SegmentPredictor
 
-def produce_depth_map(image):
-    depth_predictor = DepthPredictor()
-    depth_result = depth_predictor.predict(image)
-    return depth_result
-
-def produce_segmentation_map(image):
-    segment_predictor = SegmentPredictor()
-    sam_result = segment_predictor.predict(image)
-    return sam_result
-
-def produce_3d_reconstruction(image):
-    depth_predictor = DepthPredictor()
-    depth_result = depth_predictor.predict(image)
-    rgb_gltf_path = create_3d_obj(np.array(image), depth_result, path='./rgb.gltf')
-    return rgb_gltf_path
-
-def produce_point_cloud(depth_map, segmentation_map):
-    return point_cloud(np.array(segmentation_map), depth_map)
-
-def snap(image, depth_map, segmentation_map):
-    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 and depth_map) else None
-
-    return [image, depth_result, sam_result, rgb_gltf_path, point_cloud_fig]
-demo = gr.Interface(
-    snap,
-    inputs=[gr.Image(source="webcam", tool=None, label="Input Image", type="pil"),
-            "checkbox",
-            "checkbox"],
-    outputs=[gr.Image(label="RGB"),
-             gr.Image(label="predicted depth"),
-             gr.Image(label="predicted segmentation"),
-             gr.Model3D(label="3D mesh reconstruction - RGB",
-                        clear_color=[1.0, 1.0, 1.0, 1.0]),
-             gr.Plot()]
-)
-
-if __name__ == "__main__":
-    demo.launch()
+
+
+sam = SegmentPredictor() #service.get_sam(configs.model_type, configs.model_ckpt_path, configs.device)
+red = (255,0,0)
+blue = (0,0,255)
+
+
+block = gr.Blocks()
+with block:
+    # States
+    def point_coords_empty():
+        return []
+    def point_labels_empty():
+        return []
+    raw_image = gr.Image(type='pil', visible=False)
+    point_coords = gr.State(point_coords_empty)
+    point_labels = gr.State(point_labels_empty)
+    masks = gr.State()
+    cutout_idx = gr.State(set())
+
+    # UI
+    with gr.Column():
+        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)
+        with gr.Row():
+            with gr.Column(scale=1):
+                point_label_radio = gr.Radio(label='Point Label', choices=[1,0], value=1)
+                reset_btn = gr.Button('Reset')
+                sam_sgmt_everything_btn = gr.Button('Segment Everything!', variant = 'primary')
+                sam_encode_btn = gr.Button('Encode', variant = 'primary')
+                sam_decode_btn = gr.Button('Predict using points!')
+
+    # components
+    components = {point_coords, point_labels, raw_image, masks, cutout_idx, input_image,
+                  point_label_radio, reset_btn, sam_sgmt_everything_btn, sam_encode_btn,
+                  sam_decode_btn, masks_annotated_image}
+    # event - init coords
+    def on_reset_btn_click(raw_image):
+        return raw_image, point_coords_empty(), point_labels_empty(), None, []
+    reset_btn.click(on_reset_btn_click, [raw_image], [input_image, point_coords, point_labels], queue=False)
+
+    def on_input_image_upload(input_image):
+        # encode image on upload
+        return input_image, point_coords_empty(), point_labels_empty(), None
+    input_image.upload(on_input_image_upload, [input_image], [raw_image, point_coords, point_labels], queue=False)
+
+    # event - set coords
+    def on_input_image_select(input_image, point_coords, point_labels, point_label_radio, evt: gr.SelectData):
+        x, y = evt.index
+        color = red if point_label_radio == 0 else blue
+        img = np.array(input_image)
+        cv2.circle(img, (x, y), 5, color, -1)
+        img = Image.fromarray(img)
+        point_coords.append([x,y])
+        point_labels.append(point_label_radio)
+        return img, point_coords, point_labels
+    input_image.select(on_input_image_select, [input_image, point_coords, point_labels, point_label_radio], [input_image, point_coords, point_labels], queue=False)
+
+    # event - inference
+    def on_click_sam_encode_btn(inputs):
+        image = inputs[raw_image]
+        sam.encode(image)
+
+    def on_click_sam_dencode_btn(inputs):
+        image = inputs[raw_image]
+        generated_masks, _ = sam.cond_pred(pts=np.array(inputs[point_coords]), lbls=np.array(inputs[point_labels]))
+        annotated = (image, [(generated_masks[i], f'Mask {i}') for i in range(len(generated_masks))])
+        return {masks_annotated_image:annotated, 
+                masks: generated_masks, 
+                cutout_idx: set()}
+    sam_encode_btn.click(on_click_sam_encode_btn, components, [masks_annotated_image, masks, cutout_idx], queue=True)
+    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__':
+    block.queue()
+    block.launch()

app_legacy.py

@@ -0,0 +1,48 @@
+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
+import numpy as np
+
+def produce_depth_map(image):
+    depth_predictor = DepthPredictor()
+    depth_result = depth_predictor.predict(image)
+    return depth_result
+
+def produce_segmentation_map(image):
+    segment_predictor = SegmentPredictor()
+    sam_result = segment_predictor.predict(image)
+    return sam_result
+
+def produce_3d_reconstruction(image):
+    depth_predictor = DepthPredictor()
+    depth_result = depth_predictor.predict(image)
+    rgb_gltf_path = create_3d_obj(np.array(image), depth_result, path='./rgb.gltf')
+    return rgb_gltf_path
+
+def produce_point_cloud(depth_map, segmentation_map):
+    return point_cloud(np.array(segmentation_map), depth_map)
+
+def snap(image, depth_map, segmentation_map):
+    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 and depth_map) else None
+
+    return [image, depth_result, sam_result, rgb_gltf_path, point_cloud_fig]
+demo = gr.Interface(
+    snap,
+    inputs=[gr.Image(source="webcam", tool=None, label="Input Image", type="pil"),
+            "checkbox",
+            "checkbox"],
+    outputs=[gr.Image(label="RGB"),
+             gr.Image(label="predicted depth"),
+             gr.Image(label="predicted segmentation"),
+             gr.Model3D(label="3D mesh reconstruction - RGB",
+                        clear_color=[1.0, 1.0, 1.0, 1.0]),
+             gr.Plot()]
+)
+
+if __name__ == "__main__":
+    demo.launch()

inference.py

@@ -1,5 +1,5 @@
 from transformers import DPTImageProcessor, DPTForDepthEstimation
-from segment_anything import SamAutomaticMaskGenerator, sam_model_registry
+from segment_anything import SamAutomaticMaskGenerator, sam_model_registry, SamPredictor
 import gradio as gr
 import supervision as sv
 import torch
@@ -13,7 +13,7 @@ class DepthPredictor:
         self.feature_extractor = DPTImageProcessor.from_pretrained("Intel/dpt-large")
         self.model = DPTForDepthEstimation.from_pretrained("Intel/dpt-large")
         self.model.eval()
-
+    
     def predict(self, image):
         # prepare image for the model
         encoding = self.feature_extractor(image, return_tensors="pt")
@@ -44,9 +44,26 @@ class SegmentPredictor:
         MODEL_TYPE = "vit_b"
         checkpoint = "sam_vit_b_01ec64.pth"
         sam = sam_model_registry[MODEL_TYPE](checkpoint=checkpoint)
+        # Select device
+        self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
+        sam.to(device=self.device)
         self.mask_generator = SamAutomaticMaskGenerator(sam)
+        self.conditioned_pred = SamPredictor(sam)
     
-    def predict(self, image):
+    def encode(self, image):
+        image = np.array(image)
+        self.conditioned_pred.set_image(image)
+    
+    def cond_pred(self, pts, lbls):
+        masks, _, _ = self.conditioned_pred.predict(
+            point_coords=pts,
+            point_labels=lbls,
+            multimask_output=True
+            )
+        return masks
+
+
+    def segment_everything(self, image):
         image = np.array(image)
         sam_result = self.mask_generator.generate(image)
         mask_annotator = sv.MaskAnnotator()