| 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 |
|
|
