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"),
+ "",
+ ],
  ],
 )