Spaces:
Configuration error
Configuration error
| import torch | |
| import numpy as np | |
| from PIL import Image | |
| from torchvision import transforms | |
| from ml.model import PatchCNN | |
| from ml.dataset import CLASS_NAMES, PATCH_SIZE | |
| CHECKPOINT_PATH = "checkpoints/model.pth" | |
| CONFIDENCE_THRESHOLD = 0.5 # below this → fall back to color | |
| # ── load model once at module level ── | |
| # this means it loads when FastAPI starts, not on every request | |
| device = torch.device("cuda" if torch.cuda.is_available() else "cpu") | |
| model = PatchCNN(num_classes=7) | |
| model.load_state_dict(torch.load(CHECKPOINT_PATH, map_location=device)) | |
| model.eval() | |
| model.to(device) | |
| print(f"PatchCNN loaded on {device}") | |
| # ── same transform as training ── | |
| transform = transforms.Compose([ | |
| transforms.Resize((32, 32)), | |
| transforms.ToTensor(), | |
| transforms.Normalize( | |
| mean=[0.485, 0.456, 0.406], | |
| std=[0.229, 0.224, 0.225] | |
| ) | |
| ]) | |
| def predict_image(image): | |
| """ | |
| Takes a PIL image. | |
| Returns: | |
| label_grid — 2D list of class name strings | |
| confidence_grid — 2D list of float confidence scores | |
| """ | |
| image_np = np.array(image.convert("RGB")) | |
| H, W, _ = image_np.shape | |
| if H < PATCH_SIZE or W < PATCH_SIZE: | |
| return [], [] | |
| label_grid = [] | |
| confidence_grid = [] | |
| for y in range(0, H - PATCH_SIZE, PATCH_SIZE): # stride = PATCH_SIZE (no overlap at inference) | |
| label_row = [] | |
| conf_row = [] | |
| for x in range(0, W - PATCH_SIZE, PATCH_SIZE): | |
| # extract patch | |
| patch = image_np[y:y+PATCH_SIZE, x:x+PATCH_SIZE] | |
| patch_pil = Image.fromarray(patch) | |
| # transform → tensor | |
| tensor = transform(patch_pil).unsqueeze(0).to(device) | |
| # unsqueeze(0) adds batch dimension: (3,32,32) → (1,3,32,32) | |
| # inference | |
| with torch.no_grad(): | |
| output = model(tensor) # (1, 7) raw logits | |
| probs = torch.softmax(output, dim=1) # convert to probabilities | |
| confidence, pred = probs.max(dim=1) # get top class + its confidence | |
| label = CLASS_NAMES[pred.item()] | |
| conf = confidence.item() | |
| label_row.append(label) | |
| conf_row.append(conf) | |
| label_grid.append(label_row) | |
| confidence_grid.append(conf_row) | |
| return label_grid, confidence_grid |