Enhance accuracy

app.py

@@ -12,13 +12,13 @@ import PIL
 import torch
 from segment_anything import SamAutomaticMaskGenerator, sam_model_registry
 
-
 CHECKPOINT_PATH = os.path.join(os.path.expanduser("~"), ".cache", "SAM")
 CHECKPOINT_NAME = "sam_vit_h_4b8939.pth"
 CHECKPOINT_URL = "https://dl.fbaipublicfiles.com/segment_anything/sam_vit_h_4b8939.pth"
 MODEL_TYPE = "default"
-MAX_WIDTH = MAX_HEIGHT = 800
-THRESHOLD = 0.05
+MAX_WIDTH = MAX_HEIGHT = 1024
+TOP_K_OBJ = 100
+THRESHOLD = 0.85
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
 
@@ -55,23 +55,19 @@ def adjust_image_size(image: np.ndarray) -> np.ndarray:
 
 
 @torch.no_grad()
-def get_scores(crops: List[PIL.Image.Image], query: str) -> torch.Tensor:
+def get_score(crop: PIL.Image.Image, texts: List[str]) -> torch.Tensor:
     model, preprocess = load_clip()
-    preprocessed = [preprocess(crop) for crop in crops]
-    preprocessed = torch.stack(preprocessed).to(device)
-    token = clip.tokenize(query).to(device)
-    img_features = model.encode_image(preprocessed)
-    txt_features = model.encode_text(token)
-    img_features /= img_features.norm(dim=-1, keepdim=True)
-    txt_features /= txt_features.norm(dim=-1, keepdim=True)
-    similarity = (100 * img_features @ txt_features.T).softmax(0)
-    return similarity
+    preprocessed = preprocess(crop).unsqueeze(0).to(device)
+    tokens = clip.tokenize(texts).to(device)
+    logits_per_image, _ = model(preprocessed, tokens)
+    similarity = logits_per_image.softmax(-1).cpu()
+    return similarity[0, 0]
 
 
 def crop_image(image: np.ndarray, mask: Dict[str, Any]) -> PIL.Image.Image:
     x, y, w, h = mask["bbox"]
     masked = image * np.expand_dims(mask["segmentation"], -1)
-    crop = masked[y: y + h, x: x + w]
+    crop = masked[y : y + h, x : x + w]
     if h > w:
         top, bottom, left, right = 0, 0, (h - w) // 2, (h - w) // 2
     else:
@@ -86,11 +82,14 @@ def crop_image(image: np.ndarray, mask: Dict[str, Any]) -> PIL.Image.Image:
         cv2.BORDER_CONSTANT,
         value=(0, 0, 0),
     )
-    crop = cv2.cvtColor(crop, cv2.COLOR_BGR2RGB)
     crop = PIL.Image.fromarray(crop)
     return crop
 
 
+def get_texts(query: str) -> List[str]:
+    return [f"a picture of {query}", "a picture of background"]
+
+
 def filter_masks(
     image: np.ndarray,
     masks: List[Dict[str, Any]],
@@ -99,26 +98,19 @@ def filter_masks(
     query: str,
     clip_threshold: float,
 ) -> List[Dict[str, Any]]:
-    cropped_masks: List[PIL.Image.Image] = []
     filtered_masks: List[Dict[str, Any]] = []
 
-    for mask in masks:
+    for mask in sorted(masks, key=lambda mask: mask["area"])[-TOP_K_OBJ:]:
         if (
             mask["predicted_iou"] < predicted_iou_threshold
             or mask["stability_score"] < stability_score_threshold
             or image.shape[:2] != mask["segmentation"].shape[:2]
+            or query
+            and get_score(crop_image(image, mask), get_texts(query)) < clip_threshold
         ):
             continue
-        filtered_masks.append(mask)
-        cropped_masks.append(crop_image(image, mask))
 
-    if query and filtered_masks:
-        scores = get_scores(cropped_masks, query)
-        filtered_masks = [
-            filtered_masks[i]
-            for i, score in enumerate(scores)
-            if score > clip_threshold
-        ]
+        filtered_masks.append(mask)
 
     return filtered_masks
 
@@ -140,7 +132,7 @@ def draw_masks(
         contours, _ = cv2.findContours(
             np.uint8(mask["segmentation"]), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE
         )
-        cv2.drawContours(image, contours, -1, (255, 0, 0), 2)
+        cv2.drawContours(image, contours, -1, (0, 0, 255), 2)
     return image
 
 
@@ -152,8 +144,11 @@ def segment(
     query: str,
 ) -> PIL.ImageFile.ImageFile:
     mask_generator = load_mask_generator()
+    image = cv2.imread(image_path, cv2.IMREAD_COLOR)
+    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+
     # reduce the size to save gpu memory
-    image = adjust_image_size(cv2.imread(image_path))
+    image = adjust_image_size(image)
     masks = mask_generator.generate(image)
     masks = filter_masks(
         image,
@@ -164,7 +159,6 @@ def segment(
         clip_threshold,
     )
     image = draw_masks(image, masks)
-    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
     image = PIL.Image.fromarray(image)
     return image
 
@@ -185,31 +179,52 @@ demo = gr.Interface(
         [
             0.9,
             0.8,
-            0.15,
+            0.99,
             os.path.join(os.path.dirname(__file__), "examples/dog.jpg"),
-            "A dog",
+            "dog",
         ],
         [
             0.9,
             0.8,
-            0.001,
+            0.75,
             os.path.join(os.path.dirname(__file__), "examples/city.jpg"),
             "building",
         ],
         [
             0.9,
             0.8,
-            0.05,
+            0.998,
             os.path.join(os.path.dirname(__file__), "examples/food.jpg"),
-            "spoon",
+            "strawberry",
         ],
         [
             0.9,
             0.8,
-            0.05,
+            0.75,
             os.path.join(os.path.dirname(__file__), "examples/horse.jpg"),
             "horse",
         ],
+        [
+            0.9,
+            0.8,
+            0.99,
+            os.path.join(os.path.dirname(__file__), "examples/bears.jpg"),
+            "bear",
+        ],
+        [
+            0.9,
+            0.8,
+            0.99,
+            os.path.join(os.path.dirname(__file__), "examples/cats.jpg"),
+            "cat",
+        ],
+        [
+            0.9,
+            0.8,
+            0.99,
+            os.path.join(os.path.dirname(__file__), "examples/fish.jpg"),
+            "fish",
+        ],
     ],
 )