File size: 2,722 Bytes
7bf97d3 | 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 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 | # model_utils.py
import json
import torch
import torch.nn.functional as F
from torchvision import models, transforms
from PIL import Image
# Assumptions: ImageNet normalization, 224x224 input — change if yours differs
IMAGENET_MEAN = [0.485, 0.456, 0.406]
IMAGENET_STD = [0.229, 0.224, 0.225]
INPUT_SIZE = 224
# Preprocessing transform (resize->center crop->to tensor->normalize)
transform = transforms.Compose([
transforms.Resize(256),
transforms.CenterCrop(INPUT_SIZE),
transforms.ToTensor(),
transforms.Normalize(IMAGENET_MEAN, IMAGENET_STD)
])
def load_labels(path):
with open(path, "r", encoding="utf-8") as f:
return [line.strip() for line in f if line.strip()]
def load_remedies(path):
with open(path, "r", encoding="utf-8") as f:
return json.load(f)
def build_model(num_classes, device, checkpoint_path):
# Build MobileNetV2 with custom classifier (must match training)
model = models.mobilenet_v2(pretrained=False)
num_ftrs = model.classifier[1].in_features
model.classifier[1] = torch.nn.Linear(num_ftrs, num_classes)
# load weights
state = torch.load(checkpoint_path, map_location=device)
# if you saved state_dict only, this works:
if isinstance(state, dict) and ("state_dict" in state) and not any(k.startswith("module.") for k in state):
model.load_state_dict(state["state_dict"])
else:
try:
model.load_state_dict(state)
except Exception:
# attempt to handle possible 'module.' prefixes (from DataParallel)
new_state = {}
for k,v in state.items():
name = k.replace("module.", "") if k.startswith("module.") else k
new_state[name] = v
model.load_state_dict(new_state)
model.to(device)
model.eval()
return model
def load_model(checkpoint_path, labels_path, remedies_path, device):
labels = load_labels(labels_path)
remedies = load_remedies(remedies_path)
model = build_model(len(labels), device, checkpoint_path)
return model, labels, remedies
def predict(model, pil_image, labels, device, topk=3):
"""Return top-1 label, confidence, and topk list of (label, prob)."""
img_t = transform(pil_image).unsqueeze(0).to(device) # shape 1x3xHxW
with torch.no_grad():
outputs = model(img_t) # logits
probs = F.softmax(outputs, dim=1) # convert to probabilities
top_probs, top_idxs = probs.topk(topk, dim=1)
top_probs = top_probs.cpu().numpy()[0]
top_idxs = top_idxs.cpu().numpy()[0]
top_labels = [labels[i] for i in top_idxs]
return top_labels[0], float(top_probs[0]), list(zip(top_labels, top_probs.tolist()))
|