Create model.py

model.py

@@ -0,0 +1,20 @@
+from transformers import AutoImageProcessor, AutoModelForSemanticSegmentation
+from PIL import Image
+import torch
+import numpy as np
+
+# Load model
+processor = AutoImageProcessor.from_pretrained("nvidia/segformer-b0-finetuned-ade-512-512")
+model = AutoModelForSemanticSegmentation.from_pretrained("nvidia/segformer-b0-finetuned-ade-512-512")
+
+def predict_defect(image: Image.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()
+    
+    # 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)