spaces

app.py

@@ -6,8 +6,10 @@ import torch
 from gradio_image_prompter import ImagePrompter
 from sam2.sam2_image_predictor import SAM2ImagePredictor
 from omegaconf import OmegaConf
+from PIL import Image
+import numpy as np
 
-from objctrl_2_5d.utils.ui_utils import process_image, get_camera_pose, run_segment, run_depth, get_points, undo_points
+from objctrl_2_5d.utils.ui_utils import process_image, get_camera_pose, get_subject_points, run_depth, get_points, undo_points, mask_image
 
 
 from cameractrl.inference import get_pipeline
@@ -114,6 +116,51 @@ pipeline = get_pipeline(model_id, "unet", model_config['down_block_types'], mode
 # d_model_NK = None
 # pipeline = None
 
+### run the demo ##
+@spaces.GPU(duration=50)
+# def run_segment(segmentor):
+def segment(canvas, image, logits):
+    if logits is not None:
+        logits *=  32.0
+    _, points = get_subject_points(canvas)
+    image = np.array(image)
+
+    with torch.inference_mode(), torch.autocast("cuda", dtype=torch.bfloat16):
+        segmentor.set_image(image)
+        input_points = []
+        input_boxes = []
+        for p in points:
+            [x1, y1, _, x2, y2, _] = p
+            if x2==0 and y2==0:
+                input_points.append([x1, y1])
+            else:
+                input_boxes.append([x1, y1, x2, y2])
+        if len(input_points) == 0:
+            input_points = None
+            input_labels = None
+        else:
+            input_points = np.array(input_points)
+            input_labels = np.ones(len(input_points))
+        if len(input_boxes) == 0:
+            input_boxes = None
+        else:
+            input_boxes = np.array(input_boxes)
+        masks, _, logits = segmentor.predict(
+            point_coords=input_points,
+            point_labels=input_labels,
+            box=input_boxes,
+            multimask_output=False,
+            return_logits=True,
+            mask_input=logits,
+        )
+        mask = masks > 0
+        masked_img = mask_image(image, mask[0], color=[252, 140, 90], alpha=0.9)
+        masked_img = Image.fromarray(masked_img)
+        
+    return mask[0], masked_img, masked_img, logits / 32.0
+
+    # return segment
+
 
 # -------------- UI definition --------------
 with gr.Blocks() as demo:
@@ -270,7 +317,8 @@ with gr.Blocks() as demo:
     )
     
     select_button.click(
-        run_segment(segmentor),
+        # run_segment(segmentor),
+        segment,
         [canvas, original_image, mask_logits],
         [mask, mask_output, masked_original_image, mask_logits]
     )

objctrl_2_5d/utils/ui_utils.py

@@ -52,49 +52,6 @@ def process_image(raw_image):
 def get_subject_points(canvas):
     return canvas["image"], canvas["points"]
 
-@spaces.GPU(duration=50)
-def run_segment(segmentor):
-    def segment(canvas, image, logits):
-        if logits is not None:
-            logits *=  32.0
-        _, points = get_subject_points(canvas)
-        image = np.array(image)
-
-        with torch.inference_mode(), torch.autocast("cuda", dtype=torch.bfloat16):
-            segmentor.set_image(image)
-            input_points = []
-            input_boxes = []
-            for p in points:
-                [x1, y1, _, x2, y2, _] = p
-                if x2==0 and y2==0:
-                    input_points.append([x1, y1])
-                else:
-                    input_boxes.append([x1, y1, x2, y2])
-            if len(input_points) == 0:
-                input_points = None
-                input_labels = None
-            else:
-                input_points = np.array(input_points)
-                input_labels = np.ones(len(input_points))
-            if len(input_boxes) == 0:
-                input_boxes = None
-            else:
-                input_boxes = np.array(input_boxes)
-            masks, _, logits = segmentor.predict(
-                point_coords=input_points,
-                point_labels=input_labels,
-                box=input_boxes,
-                multimask_output=False,
-                return_logits=True,
-                mask_input=logits,
-            )
-            mask = masks > 0
-            masked_img = mask_image(image, mask[0], color=[252, 140, 90], alpha=0.9)
-            masked_img = Image.fromarray(masked_img)
-            
-        return mask[0], masked_img, masked_img, logits / 32.0
-
-    return segment
 
 def mask_image(image,
                mask,