contour instead of overlay

app.py

@@ -19,6 +19,7 @@ import random
 from functools import lru_cache
 from typing import Any, Dict, List, Optional, Tuple, Union
 
+import cv2
 import gradio as gr
 import numpy as np
 import pandas as pd
@@ -29,7 +30,6 @@ from transformers import (
     AutoImageProcessor,
     AutoModelForImageClassification,
 )
-from torchvision.utils import draw_segmentation_masks
 
 
 HF_REPO_ID = "raidium/curia"
@@ -214,15 +214,31 @@ def prepare_mask_tensor(mask: Any, height: int, width: int) -> Optional[torch.Te
     return stacked
 
 
-def apply_mask_overlay(image: np.ndarray, mask: Any) -> np.ndarray:
+def apply_contour_overlay(
+    image: np.ndarray,
+    mask: Any,
+    thickness: int = 1,
+    color: Tuple[int, int, int] = (255, 0, 0),
+) -> np.ndarray:
+    """Draw only the contours of segmentation masks instead of filled masks."""
     height, width = image.shape[:2]
     mask_tensor = prepare_mask_tensor(mask, height, width)
     if mask_tensor is None:
         return image
 
-    img_tensor = torch.from_numpy(image).permute(2, 0, 1).float() / 255.0
-    overlaid = draw_segmentation_masks(img_tensor, mask_tensor, colors=[(255, 0, 0)], alpha=0.4)
-    output = overlaid.permute(1, 2, 0).mul(255).clamp(0, 255).byte().numpy()
+    # Work on a copy of the image
+    output = image.copy()
+
+    # Process each mask separately
+    for idx in range(mask_tensor.shape[0]):
+        mask_np = mask_tensor[idx].numpy().astype(np.uint8)
+
+        # Find contours
+        contours, _ = cv2.findContours(mask_np, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+
+        # Draw contours on the image
+        cv2.drawContours(output, contours, -1, color, thickness)
+
     return output
 
 
@@ -232,7 +248,7 @@ def render_image_with_mask_info(image: np.ndarray, mask: Any) -> Tuple[np.ndarra
         return display, None
 
     try:
-        overlaid = apply_mask_overlay(display, mask)
+        overlaid = apply_contour_overlay(display, mask)
         return overlaid, ""
     except Exception:
         return display, "Mask provided but could not be visualised."
@@ -552,8 +568,8 @@ def build_demo() -> gr.Blocks:
 
             - Configure the `HF_TOKEN` secret in your Space to load private checkpoints
               and datasets from the `raidium` organisation.
-            - When masks are available in the dataset sample, they are overlaid on the
-              image for visual reference (courtesy of `torchvision.utils.draw_segmentation_masks`).
+            - When masks are available in the dataset sample, their contours are drawn on the
+              image for visual reference using OpenCV.
             - Uploaded images must be single-channel arrays. Multi-channel inputs are
               converted to grayscale automatically.
             """

requirements.txt

@@ -2,8 +2,8 @@ gradio>=4.44.0
 transformers>=4.41.0
 datasets>=2.19.0
 torch>=2.2.0
-torchvision>=0.17.0
 pandas>=2.2.0
 numpy>=1.26.0
 pillow>=10.2.0
+opencv-python>=4.8.0