File size: 1,253 Bytes
a10bf26 | 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 | import json
import torch
import torch.nn as nn
from PIL import Image, ImageOps
from torchvision import models, transforms
def build_model():
model = models.efficientnet_b0(weights=None)
model.classifier = nn.Sequential(
nn.Dropout(p=0.2),
nn.Linear(in_features=1280, out_features=1),
)
return model
def load_model(weights_path, device="cpu"):
model = build_model()
model.load_state_dict(torch.load(weights_path, map_location=device))
model.to(device)
model.eval()
return model
def predict(model, image_path, config_path, device="cpu"):
with open(config_path) as f:
cfg = json.load(f)
transform = transforms.Compose([
transforms.Resize(cfg["resize_before_crop"]),
transforms.CenterCrop(cfg["image_size"]),
transforms.ToTensor(),
transforms.Normalize(cfg["normalize_mean"], cfg["normalize_std"]),
])
img = ImageOps.exif_transpose(Image.open(image_path)).convert("RGB")
x = transform(img).unsqueeze(0).to(device)
with torch.no_grad():
prob_class1 = torch.sigmoid(model(x)).item()
idx = int(prob_class1 > 0.5)
confidence = prob_class1 if idx == 1 else 1 - prob_class1
return cfg["class_names"][idx], confidence
|