Update inference.py

inference.py

@@ -4,23 +4,83 @@ 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
+
 # Load model from HF (swap this with your own if you want)
 HF_MODEL_ID = "EPFL-ECEO/segformer-b2-finetuned-coralscapes-1024-1024"
 
 class CoralSegModel:
     def __init__(self, device=None):
         self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")
+
         self.processor = SegformerImageProcessorFast.from_pretrained(HF_MODEL_ID)
-        self.model = SegformerForSemanticSegmentation.from_pretrained(HF_MODEL_ID).to(self.device)
+
+        self.model = SegformerForSemanticSegmentation.from_pretrained(
+            HF_MODEL_ID,
+            dtype=torch.bfloat16
+        ).to(self.device)
+
         self.model.eval()
 
-        # Build a simple color palette for masks (fallback if none provided)
-        # 0..N-1 colors - here random-ish but stable
-        num_classes = self.model.config.id2label and len(self.model.config.id2label) or 40
-        rng = np.random.RandomState(0)
-        self.palette = (rng.randint(0, 255, size=(num_classes, 3))).astype(np.uint8)
+        self.palette = make_palette(index_to_name, label2color)
 
-    @spaces.GPU
+    @torch.inference_mode()
     def predict_overlay(self, frame_bgr: np.ndarray, alpha: float = 0.45) -> np.ndarray:
         """
         frame_bgr: np.ndarray HxWx3 in BGR (as read by OpenCV)
@@ -30,7 +90,7 @@ class CoralSegModel:
         rgb = frame_bgr[:, :, ::-1]
         pil = Image.fromarray(rgb)
 
-        inputs = self.processor(images=pil, return_tensors="pt", device=self.device).to(device)
+        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(
@@ -40,5 +100,4 @@ class CoralSegModel:
 
         color_mask = self.palette[pred]  # HxWx3 (RGB)
         overlay_rgb = (rgb * (1 - alpha) + color_mask * alpha).astype(np.uint8)
-        overlay_bgr = overlay_rgb[:, :, ::-1]
-        return overlay_bgr
+        return overlay_rgb