SuriRaja commited on
Commit
8470e33
·
verified ·
1 Parent(s): e6fd192

Create model.py

Browse files
Files changed (1) hide show
  1. model.py +20 -0
model.py ADDED
@@ -0,0 +1,20 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from transformers import AutoImageProcessor, AutoModelForSemanticSegmentation
2
+ from PIL import Image
3
+ import torch
4
+ import numpy as np
5
+
6
+ # Load model
7
+ processor = AutoImageProcessor.from_pretrained("nvidia/segformer-b0-finetuned-ade-512-512")
8
+ model = AutoModelForSemanticSegmentation.from_pretrained("nvidia/segformer-b0-finetuned-ade-512-512")
9
+
10
+ def predict_defect(image: Image.Image):
11
+ inputs = processor(images=image, return_tensors="pt")
12
+ with torch.no_grad():
13
+ outputs = model(**inputs)
14
+ logits = outputs.logits
15
+ segmentation = torch.argmax(logits.squeeze(), dim=0).detach().cpu().numpy()
16
+
17
+ # Convert to RGB overlay
18
+ overlay = np.zeros((segmentation.shape[0], segmentation.shape[1], 3), dtype=np.uint8)
19
+ overlay[segmentation == 12] = [255, 0, 0] # example label index for defects (adjust accordingly)
20
+ return Image.fromarray(overlay)