Update inference.py

inference.py

@@ -2,7 +2,7 @@
 import torch
 import numpy as np
 from PIL import Image
-from transformers import SegformerImageProcessor, SegformerForSemanticSegmentation
+from transformers import SegformerImageProcessorFast, SegformerForSemanticSegmentation
 
 # Load model from HF (swap this with your own if you want)
 HF_MODEL_ID = "EPFL-ECEO/segformer-b2-finetuned-coralscapes-1024-1024"
@@ -10,7 +10,7 @@ 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 = SegformerImageProcessor.from_pretrained(HF_MODEL_ID)
+        self.processor = SegformerImageProcessorFast.from_pretrained(HF_MODEL_ID)
         self.model = SegformerForSemanticSegmentation.from_pretrained(HF_MODEL_ID).to(self.device)
         self.model.eval()
 
@@ -30,7 +30,7 @@ class CoralSegModel:
         rgb = frame_bgr[:, :, ::-1]
         pil = Image.fromarray(rgb)
 
-        inputs = self.processor(images=pil, return_tensors="pt").to(self.device)
+        inputs = self.processor(images=pil, return_tensors="pt", device=self.device)
         outputs = self.model(**inputs)
         logits = outputs.logits  # [B, C, h, w]
         upsampled = torch.nn.functional.interpolate(