File size: 734 Bytes
1bc9e9e |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 |
import torch
from model import SimpleCNN
from PIL import Image
from torchvision import transforms
# Load model
device = "cuda" if torch.cuda.is_available() else "cpu"
model = SimpleCNN()
state = torch.load("pytorch_model.bin", map_location="cpu")
model.load_state_dict(state)
model.eval()
# Preprocessing
transform = transforms.Compose([
transforms.Resize((224, 224)),
transforms.ToTensor(),
])
labels = ["no_crack", "crack"]
def predict(image: Image.Image):
img = transform(image).unsqueeze(0)
with torch.no_grad():
logits = model(img)
probs = torch.softmax(logits, dim=1)[0]
idx = probs.argmax().item()
return {
"label": labels[idx],
"score": float(probs[idx])
}
|