|
|
import torch, json |
|
|
import torchvision |
|
|
from torchvision import transforms, models |
|
|
from PIL import Image |
|
|
|
|
|
def build_model(arch, dropout, width, freeze_backbone, num_classes=2): |
|
|
import torch.nn as nn |
|
|
if arch == "smallcnn": |
|
|
class SmallCNN(nn.Module): |
|
|
def __init__(self, num_classes=2, dropout=0.2, width=32): |
|
|
super().__init__() |
|
|
c = width |
|
|
self.features = nn.Sequential( |
|
|
nn.Conv2d(3, c, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2), |
|
|
nn.Conv2d(c, 2*c, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2), |
|
|
nn.Conv2d(2*c, 4*c, 3, padding=1), nn.ReLU(), nn.AdaptiveAvgPool2d(1), |
|
|
) |
|
|
self.head = nn.Sequential(nn.Flatten(), nn.Dropout(dropout), nn.Linear(4*c, num_classes)) |
|
|
def forward(self, x): return self.head(self.features(x)) |
|
|
return SmallCNN(num_classes=num_classes, dropout=dropout, width=width) |
|
|
elif arch == "resnet18": |
|
|
m = models.resnet18(weights=None) |
|
|
in_features = m.fc.in_features |
|
|
import torch.nn as nn |
|
|
m.fc = nn.Sequential(nn.Dropout(dropout), nn.Linear(in_features, num_classes)) |
|
|
return m |
|
|
elif arch == "mobilenet_v3_small": |
|
|
m = models.mobilenet_v3_small(weights=None) |
|
|
in_features = m.classifier[-1].in_features |
|
|
import torch.nn as nn |
|
|
m.classifier[-1] = nn.Linear(in_features, num_classes) |
|
|
return m |
|
|
else: |
|
|
raise ValueError("Unknown arch") |
|
|
|
|
|
def load_model(model_path="model_state.pt", config_path="config.json", device="cpu"): |
|
|
with open(config_path) as f: |
|
|
cfg = json.load(f) |
|
|
model = build_model(cfg["arch"], cfg["dropout"], cfg["width"], cfg["freeze_backbone"], cfg["num_classes"]) |
|
|
state = torch.load(model_path, map_location=device) |
|
|
model.load_state_dict(state, strict=True) |
|
|
model.to(device).eval() |
|
|
tfm = transforms.Compose([ |
|
|
transforms.Resize(int(cfg["img_size"]*1.14)), |
|
|
transforms.CenterCrop(cfg["img_size"]), |
|
|
transforms.ToTensor(), |
|
|
transforms.Normalize(mean=cfg["mean"], std=cfg["std"]), |
|
|
]) |
|
|
return model, tfm, cfg |
|
|
|
|
|
def predict_image(image_path, model, tfm, device="cpu"): |
|
|
img = Image.open(image_path).convert("RGB") |
|
|
x = tfm(img).unsqueeze(0).to(device) |
|
|
with torch.no_grad(): |
|
|
logits = model(x) |
|
|
probs = torch.softmax(logits, dim=1).cpu().numpy().ravel().tolist() |
|
|
pred = int(logits.argmax(dim=1).item()) |
|
|
return pred, probs |
|
|
|
