File size: 869 Bytes
8470e33 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 |
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)
|