Update inference.py

inference.py

@@ -1,69 +1,72 @@
 # inference.py
 import torch
+import torch.nn.functional as F
+
+import json
+import urllib.request
+import cv2
 import numpy as np
 from PIL import Image
 from transformers import SegformerImageProcessorFast, SegformerForSemanticSegmentation
 
-id2label   = {
-    '1': 'seagrass',
-    '2': 'trash',
-    '3': 'other coral dead',
-    '4': 'other coral bleached',
-    '5': 'sand',
-    '6': 'other coral alive',
-    '7': 'human',
-    '8': 'transect tools',
-    '9': 'fish',
-    '10': 'algae covered substrate',
-    '11': 'other animal',
-    '12': 'unknown hard substrate',
-    '13': 'background',
-    '14': 'dark',
-    '15': 'transect line',
-    '16': 'massive/meandering bleached',
-    '17': 'massive/meandering alive',
-    '18': 'rubble',
-    '19': 'branching bleached',
-    '20': 'branching dead',
-    '21': 'millepora',
-    '22': 'branching alive',
-    '23': 'massive/meandering dead',
-    '24': 'clam',
-    '25': 'acropora alive',
-    '26': 'sea cucumber',
-    '27': 'turbinaria',
-    '28': 'table acropora alive',
-    '29': 'sponge',
-    '30': 'anemone',
-    '31': 'pocillopora alive',
-    '32': 'table acropora dead',
-    '33': 'meandering bleached',
-    '34': 'stylophora alive',
-    '35': 'sea urchin',
-    '36': 'meandering alive',
-    '37': 'meandering dead',
-    '38': 'crown of thorn',
-    '39': 'dead clam'
- }
-
-label2color= {'human': [255, 0, 0], 'background': [29, 162, 216], 'fish': [255, 255, 0], 'sand': [194, 178, 128], 'rubble': [161, 153, 128], 'unknown hard substrate': [125, 125, 125], 'algae covered substrate': [125, 163, 125], 'dark': [31, 31, 31], 'branching bleached': [252, 231, 240], 'branching dead': [123, 50, 86], 'branching alive': [226, 91, 157], 'stylophora alive': [255, 111, 194], 'pocillopora alive': [255, 146, 150], 'acropora alive': [236, 128, 255], 'table acropora alive': [189, 119, 255], 'table acropora dead': [85, 53, 116], 'millepora': [244, 150, 115], 'turbinaria': [228, 255, 119], 'other coral bleached': [250, 224, 225], 'other coral dead': [114, 60, 61], 'other coral alive': [224, 118, 119], 'massive/meandering alive': [236, 150, 21], 'massive/meandering dead': [134, 86, 18], 'massive/meandering bleached': [255, 248, 228], 'meandering alive': [230, 193, 0], 'meandering dead': [119, 100, 14], 'meandering bleached': [251, 243, 216], 'transect line': [0, 255, 0], 'transect tools': [8, 205, 12], 'sea urchin': [0, 142, 255], 'sea cucumber': [0, 231, 255], 'anemone': [0, 255, 189], 'sponge': [240, 80, 80], 'clam': [189, 255, 234], 'other animal': [0, 255, 255], 'trash': [255, 0, 134], 'seagrass': [125, 222, 125], 'crown of thorn': [179, 245, 234], 'dead clam': [89, 155, 134]}                        # {'seagrass':[R,G,B],...}
-
-# Helper: build a palette aligned to class indices
-# We assume your model outputs class ids in [0..38].
-# We map index 0-> id "1", index 1-> id "2", ..., index 38-> id "39".
-# If your model uses a *different* order, define a custom `index_to_id` list accordingly.
-index_to_id = [str(i) for i in range(1, 40)]  # ["1","2",...,"39"]
-index_to_name = [id2label[i] for i in index_to_id]
-
-def make_palette(index_to_name, label2color):
-    palette = np.zeros((len(index_to_name), 3), dtype=np.uint8)
-    for k, name in enumerate(index_to_name):
-        rgb = label2color.get(name, [0, 0, 0])
-        palette[k] = np.array(rgb, dtype=np.uint8)
-    return palette
+id2label = json.load(urllib.request.urlopen(
+    "https://huggingface.co/datasets/EPFL-ECEO/coralscapes/resolve/main/id2label.json"))
+label2color = json.load(urllib.request.urlopen(
+    "https://huggingface.co/datasets/EPFL-ECEO/coralscapes/resolve/main/label2color.json"))
 
 # Load model from HF (swap this with your own if you want)
-HF_MODEL_ID = "EPFL-ECEO/segformer-b2-finetuned-coralscapes-1024-1024"
+HF_MODEL_ID = "EPFL-ECEO/segformer-b5-finetuned-coralscapes-1024-1024"
+
+def create_segmentation_overlay(pred, id2label, label2color, image, alpha=0.25):
+    """
+    Colorizes the segmentation prediction and creates an overlay image.
+
+    Args:
+        pred: The segmentation prediction (numpy array).
+        id2label: Dictionary mapping class IDs to labels.
+        label2color: Dictionary mapping labels to colors.
+        image: The original PIL Image.
+
+    Returns:
+        A PIL Image representing the overlay of the original image and the colorized segmentation mask.
+    """
+    H, W = pred.shape
+    rgb = np.zeros((H, W, 3), dtype=np.uint8)
+
+    # Get unique class IDs present in the prediction
+    unique_classes = np.unique(pred)
+
+    # Create a mapping from class ID to color
+    id2color = {int(id): label2color[label] for id, label in id2label.items()}
+
+    # Define a default color for unknown classes (e.g., black)
+    default_color = [0, 0, 0]
+
+    # Iterate through unique class IDs and colorize the image
+    for class_id in unique_classes:
+        # Get the color for the current class ID, use default_color if not found
+        rgb_c = id2color.get(int(class_id), default_color)
+        # Assign the color to the pixels with the current class ID
+        rgb[pred == class_id] = rgb_c
+
+    mask_rgb = Image.fromarray(rgb)
+
+    # 4) Alpha overlay
+    overlay = Image.blend(image.convert("RGBA"), mask_rgb.convert("RGBA"), alpha=alpha)
+
+    return overlay
+
+def resize_image(image, target_size=1024):
+    """
+    Used to resize the image such that the smaller side equals 1024
+    """
+    h_img, w_img = image.size
+    if h_img < w_img:
+        new_h, new_w = target_size, int(w_img * (target_size / h_img))
+    else:
+        new_h, new_w  = int(h_img * (target_size / w_img)), target_size
+    resized_img = image.resize((new_h, new_w))
+    return resized_img
 
 class CoralSegModel:
     def __init__(self, device=None):
@@ -78,26 +81,71 @@ class CoralSegModel:
 
         self.model.eval()
 
-        self.palette = make_palette(index_to_name, label2color)
+    @spaces.GPU
+    def segment_image(self, image, preprocessor, model, crop_size = (1024, 1024), num_classes = 40) -> np.ndarray:
+        """
+        Finds an optimal stride based on the image size and aspect ratio to create
+        overlapping sliding windows of size 1024x1024 which are then fed into the model.
+        """
+        h_crop, w_crop = crop_size
+
+        img = torch.Tensor(np.array(resize_image(image, target_size=1024)).transpose(2, 0, 1)).unsqueeze(0)
+        img = img.to(self.device, torch.bfloat16)
+        batch_size, _, h_img, w_img = img.size()
+
+        h_grids = int(np.round(3/2*h_img/h_crop)) if h_img > h_crop else 1
+        w_grids = int(np.round(3/2*w_img/w_crop)) if w_img > w_crop else 1
+
+        h_stride = int((h_img - h_crop + h_grids -1)/(h_grids -1)) if h_grids > 1 else h_crop
+        w_stride = int((w_img - w_crop + w_grids -1)/(w_grids -1)) if w_grids > 1 else w_crop
+
+        preds = img.new_zeros((batch_size, num_classes, h_img, w_img))
+        count_mat = img.new_zeros((batch_size, 1, h_img, w_img))
 
-    @torch.inference_mode()
-    def predict_overlay(self, frame_bgr: np.ndarray, alpha: float = 0.45) -> np.ndarray:
+        for h_idx in range(h_grids):
+            for w_idx in range(w_grids):
+                y1 = h_idx * h_stride
+                x1 = w_idx * w_stride
+                y2 = min(y1 + h_crop, h_img)
+                x2 = min(x1 + w_crop, w_img)
+                y1 = max(y2 - h_crop, 0)
+                x1 = max(x2 - w_crop, 0)
+                crop_img = img[:, :, y1:y2, x1:x2]
+                with torch.no_grad():
+                    if(preprocessor):
+                        inputs = preprocessor(crop_img, return_tensors = "pt", device=self.device)
+                        inputs["pixel_values"] = inputs["pixel_values"].to(self.device, torch.bfloat16)
+                    else:
+                        inputs = crop_img.to(self.device, torch.bfloat16)
+                    outputs = model.to(self.device)(**inputs)
+
+                resized_logits = F.interpolate(
+                    outputs.logits[0].unsqueeze(dim=0), size=crop_img.shape[-2:], mode="bilinear", align_corners=False
+                )
+                preds += F.pad(resized_logits,
+                                (int(x1), int(preds.shape[3] - x2), int(y1),
+                                int(preds.shape[2] - y2)))
+                count_mat[:, :, y1:y2, x1:x2] += 1
+
+        assert (count_mat == 0).sum() == 0
+        preds = preds / count_mat
+        preds = preds.argmax(dim=1)
+        preds = F.interpolate(preds.unsqueeze(0).type(torch.uint8), size=image.size[::-1], mode='nearest')
+        label_pred = preds.squeeze().cpu().numpy()
+        return label_pred
+
+    @spaces.GPU
+    def predict_map_and_overlay(self, frame_bgr: np.ndarray):
         """
-        frame_bgr: np.ndarray HxWx3 in BGR (as read by OpenCV)
-        returns: np.ndarray HxWx3 in BGR (overlay)
+        Returns:
+          pred_map:  HxW (uint8/int) with class indices in [0..C-1]
+          overlay:   HxWx3 RGB uint8 (blended color mask over original)
+          rgb:       HxWx3 RGB uint8 original frame (for AnnotatedImage base)
         """
-        # Convert BGR -> RGB PIL
-        rgb = frame_bgr[:, :, ::-1]
-        pil = Image.fromarray(rgb)
+        rgb = frame_bgr
 
-        inputs = self.processor(images=pil, return_tensors="pt", device=self.device).to(self.device, torch.bfloat16)
-        outputs = self.model(**inputs)
-        logits = outputs.logits  # [B, C, h, w]
-        upsampled = torch.nn.functional.interpolate(
-            logits, size=pil.size[::-1], mode="bilinear", align_corners=False
-        )
-        pred = upsampled.argmax(dim=1)[0].detach().cpu().numpy().astype(np.uint8)  # HxW
-
-        color_mask = self.palette[pred]  # HxWx3 (RGB)
-        overlay_rgb = (rgb * (1 - alpha) + color_mask * alpha).astype(np.uint8)
-        return overlay_rgb
+        pil = Image.fromarray(rgb)
+        pred = self.segment_image(pil, self.processor, self.model)
+        overlay_rgb = create_segmentation_overlay(pred, id2label, label2color, pil, 0.5)
+        
+        return pred, overlay_rgb, rgb