Update model.py

model.py

@@ -13,8 +13,11 @@ def predict_defect(image: Image.Image):
         outputs = model(**inputs)
     logits = outputs.logits
     segmentation = torch.argmax(logits.squeeze(), dim=0).detach().cpu().numpy()
-    
-    # Convert to RGB overlay
-    overlay = np.zeros((segmentation.shape[0], segmentation.shape[1], 3), dtype=np.uint8)
-    overlay[segmentation == 12] = [255, 0, 0]  # example label index for defects (adjust accordingly)
-    return Image.fromarray(overlay)
+
+    # Overlay on original image
+    original = np.array(image).copy()
+    mask = (segmentation == 12)  # Replace 12 with correct defect label
+    original[mask] = [255, 0, 0]  # Red highlight for defects
+
+    return Image.fromarray(original)
+