Update model.py

model.py

@@ -1,32 +1,28 @@
 from transformers import AutoImageProcessor, AutoModelForSemanticSegmentation
+from PIL import Image
 import torch
 import numpy as np
 import cv2
-from PIL import Image
 
-# Load model
 processor = AutoImageProcessor.from_pretrained("nvidia/segformer-b0-finetuned-ade-512-512")
 model = AutoModelForSemanticSegmentation.from_pretrained("nvidia/segformer-b0-finetuned-ade-512-512")
 
-# SegFormer ADE20K label IDs: road is class 0 or 8 depending on the mapping
-ROAD_LABELS = [0, 8]  # adjust if needed based on actual mapping
+ROAD_LABELS = [0, 8]  # class indices to consider as road
 
 def predict_defect(image: Image.Image):
     original = np.array(image)
     inputs = processor(images=image, return_tensors="pt")
     with torch.no_grad():
         outputs = model(**inputs)
-
     logits = outputs.logits
     segmentation = torch.argmax(logits.squeeze(), dim=0).detach().cpu().numpy()
-
-    # Resize mask to match original image size
     resized_mask = cv2.resize(segmentation.astype(np.uint8), (original.shape[1], original.shape[0]), interpolation=cv2.INTER_NEAREST)
 
-    # Highlight anything that's NOT road
-    mask = ~np.isin(resized_mask, ROAD_LABELS)
+    # Mark only suspicious areas (non-road) unless it's oversegmenting
+    defect_mask = ~np.isin(resized_mask, ROAD_LABELS)
+    if np.sum(defect_mask) / resized_mask.size > 0.4:
+        defect_mask[:] = False
 
     overlay = original.copy()
-    overlay[mask] = [255, 0, 0]  # red highlight for all non-road anomalies
-
-    return Image.fromarray(overlay)
+    overlay[defect_mask] = [255, 0, 0]
+    return Image.fromarray(overlay)