Add threshold

app.py

@@ -1,18 +1,19 @@
 import os
-import PIL
 from functools import lru_cache
-
 from random import randint
-import gradio as gr
+from typing import Dict, List
+
 import cv2
-import torch
+import gradio as gr
 import numpy as np
-from segment_anything import sam_model_registry, SamAutomaticMaskGenerator
-from typing import List
+import PIL
+import torch
+from segment_anything import SamAutomaticMaskGenerator, sam_model_registry
 
 CHECKPOINT_PATH = "sam_vit_h_4b8939.pth"
 MODEL_TYPE = "default"
 MAX_WIDTH = MAX_HEIGHT = 800
+THRESHOLD = 0.05
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
 
@@ -32,37 +33,63 @@ def adjust_image_size(image: np.ndarray) -> np.ndarray:
         if width > MAX_WIDTH:
             height, width = int(MAX_WIDTH / width * height), MAX_WIDTH
     image = cv2.resize(image, (width, height))
-    print(image.shape)
     return image
 
 
+def filter_masks(
+    masks: List[Dict[str, np.ndarray]],
+    predicted_iou_threshold: float,
+    stability_score_threshold: float,
+    query: str,
+    clip_threshold: float,
+) -> List[np.ndarray]:
+    filtered_masks: List[Dict[str, np.ndarray]] = []
+    for mask in masks:
+        if (
+            mask["predicted_iou"] < predicted_iou_threshold
+            or mask["stability_score"] < stability_score_threshold
+        ):
+            continue
+        filtered_masks.append(mask)
+
+    return [mask["segmentation"] for mask in filtered_masks]
+
+
 def draw_masks(
     image: np.ndarray, masks: List[np.ndarray], alpha: float = 0.7
 ) -> np.ndarray:
     for mask in masks:
         color = [randint(127, 255) for _ in range(3)]
-        segmentation = mask["segmentation"]
 
         # draw mask overlay
-        colored_seg = np.expand_dims(segmentation, 0).repeat(3, axis=0)
-        colored_seg = np.moveaxis(colored_seg, 0, -1)
-        masked = np.ma.MaskedArray(image, mask=colored_seg, fill_value=color)
+        colored_mask = np.expand_dims(mask, 0).repeat(3, axis=0)
+        colored_mask = np.moveaxis(colored_mask, 0, -1)
+        masked = np.ma.MaskedArray(image, mask=colored_mask, fill_value=color)
         image_overlay = masked.filled()
         image = cv2.addWeighted(image, 1 - alpha, image_overlay, alpha, 0)
 
         # draw contour
         contours, _ = cv2.findContours(
-            np.uint8(segmentation), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE
+            np.uint8(mask), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE
         )
         cv2.drawContours(image, contours, -1, (255, 0, 0), 2)
     return image
 
 
-def segment(image_path: str, query: str) -> PIL.ImageFile.ImageFile:
+def segment(
+    predicted_iou_threshold: float,
+    stability_score_threshold: float,
+    clip_threshold: float,
+    image_path: str,
+    query: str,
+) -> PIL.ImageFile.ImageFile:
     mask_generator = load_mask_generator()
     # reduce the size to save gpu memory
     image = adjust_image_size(cv2.imread(image_path))
     masks = mask_generator.generate(image)
+    masks = filter_masks(
+        masks, predicted_iou_threshold, stability_score_threshold, query, clip_threshold
+    )
     image = draw_masks(image, masks)
     image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
     image = PIL.Image.fromarray(np.uint8(image)).convert("RGB")
@@ -71,15 +98,45 @@ def segment(image_path: str, query: str) -> PIL.ImageFile.ImageFile:
 
 demo = gr.Interface(
     fn=segment,
-    inputs=[gr.Image(type="filepath"), "text"],
+    inputs=[
+        gr.Slider(0, 1, value=0.9, label="predicted_iou_threshold"),
+        gr.Slider(0, 1, value=0.8, label="stability_score_threshold"),
+        gr.Slider(0, 1, value=0.05, label="clip_threshold"),
+        gr.Image(type="filepath"),
+        "text",
+    ],
     outputs="image",
     allow_flagging="never",
     title="Segment Anything with CLIP",
     examples=[
-        [os.path.join(os.path.dirname(__file__), "examples/dog.jpg"), ""],
-        [os.path.join(os.path.dirname(__file__), "examples/city.jpg"), ""],
-        [os.path.join(os.path.dirname(__file__), "examples/food.jpg"), ""],
-        [os.path.join(os.path.dirname(__file__), "examples/horse.jpg"), ""],
+        [
+            0.9,
+            0.8,
+            0.05,
+            os.path.join(os.path.dirname(__file__), "examples/dog.jpg"),
+            "",
+        ],
+        [
+            0.9,
+            0.8,
+            0.05,
+            os.path.join(os.path.dirname(__file__), "examples/city.jpg"),
+            "",
+        ],
+        [
+            0.9,
+            0.8,
+            0.05,
+            os.path.join(os.path.dirname(__file__), "examples/food.jpg"),
+            "",
+        ],
+        [
+            0.9,
+            0.8,
+            0.05,
+            os.path.join(os.path.dirname(__file__), "examples/horse.jpg"),
+            "",
+        ],
     ],
 )