Added 'Visualizer' class and mask refinement to utils.py

This commit includes the addition of the 'Visualizer' class in utils.py, which provides several methods to annotate images with bounding boxes, masks, polygons, and labels. It also included new utility functions for refining and filtering masks based on their relative area.

app.py

@@ -8,10 +8,11 @@ import supervision as sv
 
 from typing import List
 from segment_anything import sam_model_registry, SamAutomaticMaskGenerator
-from utils import refine_mask
+from utils import postprocess_masks, Visualizer
 
 HOME = os.getenv("HOME")
 DEVICE = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
+MINIMUM_AREA_THRESHOLD = 0.01
 
 SAM_CHECKPOINT = os.path.join(HOME, "app/weights/sam_vit_h_4b8939.pth")
 # SAM_CHECKPOINT = "weights/sam_vit_h_4b8939.pth"
@@ -27,19 +28,34 @@ MARKDOWN = """
 </h1>
 """
 
-sam = sam_model_registry[SAM_MODEL_TYPE](checkpoint=SAM_CHECKPOINT).to(device=DEVICE)
-mask_generator = SamAutomaticMaskGenerator(sam)
+SAM = sam_model_registry[SAM_MODEL_TYPE](checkpoint=SAM_CHECKPOINT).to(device=DEVICE)
+VISUALIZER = Visualizer()
 
 
 def inference(image: np.ndarray, annotation_mode: List[str]) -> np.ndarray:
-    return image
+    mask_generator = SamAutomaticMaskGenerator(SAM)
+    result = mask_generator.generate(image=image)
+    detections = sv.Detections.from_sam(result)
+    detections = postprocess_masks(
+        detections=detections,
+        area_threshold=MINIMUM_AREA_THRESHOLD)
+    bgr_image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
+    annotated_image = VISUALIZER.visualize(
+        image=bgr_image,
+        detections=detections,
+        with_box="Box" in annotation_mode,
+        with_mask="Mask" in annotation_mode,
+        with_polygon="Polygon" in annotation_mode,
+        with_label="Mark" in annotation_mode)
+    return cv2.cvtColor(annotated_image, cv2.COLOR_BGR2RGB)
 
 
 image_input = gr.Image(
     label="Input",
-    type="numpy")
+    type="numpy",
+    height=512)
 checkbox_annotation_mode = gr.CheckboxGroup(
-    choices=["Mark", "Mask", "Box"],
+    choices=["Mark", "Polygon", "Mask", "Box"],
     value=['Mark'],
     label="Annotation Mode")
 image_output = gr.Image(

utils.py

@@ -1,6 +1,55 @@
 import cv2
 
 import numpy as np
+import supervision as sv
+
+
+class Visualizer:
+
+    def __init__(
+        self,
+        line_thickness: int = 2,
+        mask_opacity: float = 0.1,
+        text_scale: float = 0.5
+    ) -> None:
+        self.box_annotator = sv.BoundingBoxAnnotator(
+            color_lookup=sv.ColorLookup.INDEX,
+            thickness=line_thickness)
+        self.mask_annotator = sv.MaskAnnotator(
+            color_lookup=sv.ColorLookup.INDEX,
+            opacity=mask_opacity)
+        self.polygon_annotator = sv.PolygonAnnotator(
+            color_lookup=sv.ColorLookup.INDEX,
+            thickness=line_thickness)
+        self.label_annotator = sv.LabelAnnotator(
+            color_lookup=sv.ColorLookup.INDEX,
+            text_position=sv.Position.CENTER_OF_MASS,
+            text_scale=text_scale)
+
+    def visualize(
+        self,
+        image: np.ndarray,
+        detections: sv.Detections,
+        with_box: bool,
+        with_mask: bool,
+        with_polygon: bool,
+        with_label: bool
+    ) -> np.ndarray:
+        annotated_image = image.copy()
+        if with_box:
+            annotated_image = self.box_annotator.annotate(
+                scene=annotated_image, detections=detections)
+        if with_mask:
+            annotated_image = self.mask_annotator.annotate(
+                scene=annotated_image, detections=detections)
+        if with_polygon:
+            annotated_image = self.polygon_annotator.annotate(
+                scene=annotated_image, detections=detections)
+        if with_label:
+            labels = list(map(str, range(len(detections))))
+            annotated_image = self.label_annotator.annotate(
+                scene=annotated_image, detections=detections, labels=labels)
+        return annotated_image
 
 
 def refine_mask(
@@ -36,4 +85,71 @@ def refine_mask(
                 mask, [contour], -1, (0 if mode == 'islands' else 255), -1
             )
 
-    return np.where(mask > 0, 1, 0)
+    return np.where(mask > 0, 1, 0).astype(bool)
+
+
+def filter_masks_by_relative_area(
+    masks: np.ndarray,
+    min_relative_area: float = 0.02,
+    max_relative_area: float = 1.0
+) -> np.ndarray:
+    """
+    Filters out masks based on their relative area.
+
+    Parameters:
+        masks (np.ndarray): A 3D numpy array where each slice along the third dimension
+            represents a mask.
+        min_relative_area (float): Minimum relative area threshold for keeping a mask.
+        max_relative_area (float): Maximum relative area threshold for keeping a mask.
+
+    Returns:
+        np.ndarray: A 3D numpy array of filtered masks.
+    """
+    mask_areas = masks.sum(axis=(1, 2))
+    total_area = masks.shape[1] * masks.shape[2]
+    relative_areas = mask_areas / total_area
+    min_area_filter = relative_areas >= min_relative_area
+    max_area_filter = relative_areas <= max_relative_area
+    return masks[min_area_filter & max_area_filter]
+
+
+def postprocess_masks(
+    detections: sv.Detections,
+    area_threshold: float = 0.02,
+    min_relative_area: float = 0.02,
+    max_relative_area: float = 1.0
+) -> sv.Detections:
+    """
+    Post-processes the masks of detection objects by removing small islands and filling
+    small holes.
+
+    Parameters:
+        detections (sv.Detections): Detection objects to be filtered.
+        area_threshold (float): Threshold for relative area to remove or fill features.
+        min_relative_area (float): Minimum relative area threshold for detections.
+        max_relative_area (float): Maximum relative area threshold for detections.
+
+    Returns:
+        np.ndarray: Post-processed masks.
+    """
+    masks = detections.mask.copy()
+    for i in range(len(masks)):
+        masks[i] = refine_mask(
+            mask=masks[i],
+            area_threshold=area_threshold,
+            mode='islands'
+        )
+        masks[i] = refine_mask(
+            mask=masks[i],
+            area_threshold=area_threshold,
+            mode='holes'
+        )
+    masks = filter_masks_by_relative_area(
+        masks=masks,
+        min_relative_area=min_relative_area,
+        max_relative_area=max_relative_area)
+
+    return sv.Detections(
+        xyxy=sv.mask_to_xyxy(masks),
+        mask=masks
+    )