Update model.py

model.py

@@ -4,29 +4,27 @@ import torch
 import numpy as np
 import cv2
 
-# Load model and processor
-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")
 
 def predict_defect(image: Image.Image):
-    # Convert image to numpy array to keep original size
     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 segmentation mask to match original image size
+
+    # Resize mask to original image size
     resized_mask = cv2.resize(segmentation.astype(np.uint8), (original.shape[1], original.shape[0]), interpolation=cv2.INTER_NEAREST)
-    
-    # For example, highlight class index 12 (customize as needed)
-    mask = (resized_mask == 12)
-    
-    # Overlay red color on defect areas
-    overlay = original.copy()
-    overlay[mask] = [255, 0, 0]  # Red for defects
 
+    # 📌 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)
+
+    # Overlay red where defect detected
+    overlay = original.copy()
+    overlay[mask] = [255, 0, 0]
     return Image.fromarray(overlay)