| 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]) | |
| } | |