Update model.py

model.py

@@ -1,30 +1,32 @@
 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")
 
-# 🔁 Replace this with a model trained for road defects
-processor = AutoImageProcessor.from_pretrained("segments/DeepLabV3")
-model = AutoModelForSemanticSegmentation.from_pretrained("segments/DeepLabV3")
+# 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
 
 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 original image size
+    # 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)
 
-    # 📌 NOTE: Update the label index below based on your dataset
-    road_defect_label_index = 1  # Assume 1 represents cracks/potholes
-    mask = (resized_mask == road_defect_label_index)
+    # Highlight anything that's NOT road
+    mask = ~np.isin(resized_mask, ROAD_LABELS)
 
-    # Overlay red where defect detected
     overlay = original.copy()
-    overlay[mask] = [255, 0, 0]
+    overlay[mask] = [255, 0, 0]  # red highlight for all non-road anomalies
+
     return Image.fromarray(overlay)