Add refine_mask utility and update Dockerfile and app.py

Introduced a new utility function 'refine_mask' in utils.py for refining masks by removing small islands or filling small holes based on a given area threshold.

Dockerfile

@@ -42,6 +42,7 @@ RUN mkdir -p $HOME/app/weights
 RUN wget -c -O $HOME/app/weights/sam_vit_h_4b8939.pth https://dl.fbaipublicfiles.com/segment_anything/sam_vit_h_4b8939.pth
 
 COPY app.py .
+COPY utils.py .
 
 RUN find $HOME/app
 

app.py

@@ -1,4 +1,5 @@
 import os
+import cv2
 import torch
 
 import gradio as gr
@@ -7,7 +8,7 @@ import supervision as sv
 
 from typing import List
 from segment_anything import sam_model_registry, SamAutomaticMaskGenerator
-
+from utils import refine_mask
 
 HOME = os.getenv("HOME")
 DEVICE = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')

utils.py

@@ -0,0 +1,39 @@
+import cv2
+
+import numpy as np
+
+
+def refine_mask(
+    mask: np.ndarray,
+    area_threshold: float,
+    mode: str = 'islands'
+) -> np.ndarray:
+    """
+    Refines a mask by removing small islands or filling small holes based on area
+    threshold.
+
+    Parameters:
+        mask (np.ndarray): Input binary mask.
+        area_threshold (float): Threshold for relative area to remove or fill features.
+        mode (str): Operation mode ('islands' for removing islands, 'holes' for filling
+                    holes).
+
+    Returns:
+        np.ndarray: Refined binary mask.
+    """
+    mask = np.uint8(mask * 255)
+    operation = cv2.RETR_EXTERNAL if mode == 'islands' else cv2.RETR_CCOMP
+    contours, _ = cv2.findContours(
+        mask, operation, cv2.CHAIN_APPROX_SIMPLE
+    )
+    total_area = cv2.countNonZero(mask) if mode == 'islands' else mask.size
+
+    for contour in contours:
+        area = cv2.contourArea(contour)
+        relative_area = area / total_area
+        if relative_area < area_threshold:
+            cv2.drawContours(
+                mask, [contour], -1, (0 if mode == 'islands' else 255), -1
+            )
+
+    return np.where(mask > 0, 1, 0)