Add functionality for interactive mask generation

Dockerfile

@@ -31,7 +31,7 @@ WORKDIR $HOME/app
 RUN pip install torch==2.0.1+cu117 torchvision==0.15.2+cu117 -f https://download.pytorch.org/whl/torch_stable.html
 
 # Install dependencies
-RUN pip install --no-cache-dir gradio==4.5.0 opencv-python supervision==0.17.0rc3 \
+RUN pip install --no-cache-dir gradio==3.50.2 opencv-python supervision==0.17.0rc3 \
     pillow requests
 
 # Install SAM and Detectron2
@@ -45,6 +45,7 @@ RUN wget -c -O $HOME/app/weights/sam_vit_h_4b8939.pth https://dl.fbaipublicfiles
 COPY app.py .
 COPY utils.py .
 COPY gpt4v.py .
+COPY sam_utils.py .
 
 RUN find $HOME/app
 

app.py

@@ -1,15 +1,16 @@
 import os
-import cv2
-import torch
+from typing import List, Dict
 
+import cv2
 import gradio as gr
 import numpy as np
 import supervision as sv
-
-from typing import List
+import torch
 from segment_anything import sam_model_registry, SamAutomaticMaskGenerator
-from utils import postprocess_masks, Visualizer
+
 from gpt4v import prompt_image
+from utils import postprocess_masks, Visualizer
+from sam_utils import sam_interactive_inference
 
 HOME = os.getenv("HOME")
 DEVICE = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
@@ -32,24 +33,33 @@ MARKDOWN = """
 
 - [ ] Support for alphabetic labels
 - [ ] Support for Semantic-SAM (multi-level)
-- [ ] Support for interactive mode
 - [ ] Support for result highlighting
+- [ ] Support for mask filtering based on granularity
 """
 
 SAM = sam_model_registry[SAM_MODEL_TYPE](checkpoint=SAM_CHECKPOINT).to(device=DEVICE)
 
 
 def inference(
-    image: np.ndarray,
+    image_and_mask: Dict[str, np.ndarray],
     annotation_mode: List[str],
     mask_alpha: float
 ) -> np.ndarray:
+    image = image_and_mask['image']
+    mask = cv2.cvtColor(image_and_mask['mask'], cv2.COLOR_RGB2GRAY)
+    is_interactive = not np.all(mask == 0)
     visualizer = Visualizer(mask_opacity=mask_alpha)
-    mask_generator = SamAutomaticMaskGenerator(SAM)
-    result = mask_generator.generate(image=image)
-    detections = sv.Detections.from_sam(result)
-    detections = postprocess_masks(
-        detections=detections)
+    if is_interactive:
+        detections = sam_interactive_inference(
+            image=image,
+            mask=mask,
+            model=SAM)
+    else:
+        mask_generator = SamAutomaticMaskGenerator(SAM)
+        result = mask_generator.generate(image=image)
+        detections = sv.Detections.from_sam(result)
+        detections = postprocess_masks(
+            detections=detections)
     bgr_image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
     annotated_image = visualizer.visualize(
         image=bgr_image,
@@ -76,7 +86,12 @@ def prompt(message, history, image: np.ndarray, api_key: str) -> str:
 image_input = gr.Image(
     label="Input",
     type="numpy",
-    height=512)
+    height=512,
+    tool="sketch",
+    interactive=True,
+    brush_radius=20.0,
+    brush_color="#FFFFFF"
+)
 checkbox_annotation_mode = gr.CheckboxGroup(
     choices=["Mark", "Polygon", "Mask", "Box"],
     value=['Mark'],

sam_utils.py

@@ -0,0 +1,30 @@
+import numpy as np
+import supervision as sv
+
+from segment_anything.modeling.sam import Sam
+from segment_anything import SamPredictor
+
+
+def sam_interactive_inference(
+    image: np.ndarray,
+    mask: np.ndarray,
+    model: Sam
+) -> sv.Detections:
+    predictor = SamPredictor(model)
+    predictor.set_image(image)
+    masks = []
+    for polygon in sv.mask_to_polygons(mask.astype(bool)):
+        random_point_indexes = np.random.choice(polygon.shape[0], size=5, replace=True)
+        input_point = polygon[random_point_indexes]
+        input_label = np.ones(5)
+        mask = predictor.predict(
+            point_coords=input_point,
+            point_labels=input_label,
+            multimask_output=False,
+        )[0][0]
+        masks.append(mask)
+    masks = np.array(masks, dtype=bool)
+    return sv.Detections(
+        xyxy=sv.mask_to_xyxy(masks),
+        mask=masks
+    )