File size: 2,793 Bytes
0362b22 |
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 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 |
import numpy as np
import pandas as pd
import torch
import torch.nn as nn
import torchvision.models as models
import torchvision.transforms as transforms
import plotly.express as px
from huggingface_hub import hf_hub_download
from config import app_config
### Reproducibility
np.random.seed(42)
torch.manual_seed(42)
torch.cuda.manual_seed_all(42)
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False
def load_model():
### useful variable that tells us whether we should use the GPU
use_cuda = torch.cuda.is_available()
### Our model is based on VGG19 with classifier head replaced with custom layers
### Load the pretrained VGG19 model
model_transfer = models.vgg19(weights="DEFAULT")
### Instantiate a custom classifier for our 50 classes and use it with pretrained model
clf = nn.Sequential(
nn.Linear(in_features=25088, out_features=4096),
nn.ReLU(inplace=True),
nn.Dropout(p=0.5),
nn.Linear(in_features=4096, out_features=1024),
nn.ReLU(inplace=True),
nn.Dropout(p=0.5),
nn.Linear(in_features=1024, out_features=50),
nn.LogSoftmax(dim=1),
)
model_transfer.classifier = clf
if use_cuda:
model_transfer = model_transfer.cuda()
### Download the saved fine-tuned model weights and load it
model_weights = hf_hub_download(
repo_id=app_config.hf_repo_id, filename=app_config.hf_weights_file
)
if not use_cuda:
state_dict = torch.load(model_weights, map_location=torch.device("cpu"))
else:
state_dict = torch.load(model_weights)
model_transfer.load_state_dict(state_dict)
return model_transfer
def predict(img, k):
model = app_config.model
classes = app_config.classes
mean = torch.tensor(app_config.mean)
std = torch.tensor(app_config.std)
# Prepare image for model prediction
transform = transforms.Compose(
[
transforms.Resize((128, 128)),
transforms.ToTensor(),
transforms.Normalize(mean, std),
]
)
img = transform(img).unsqueeze(dim=0)
# Predict
model.eval()
use_cuda = torch.cuda.is_available()
if use_cuda:
img = img.cuda()
pred = model(img) # LogSoftmax output
pred_proba = torch.exp(pred) # Prediction probabilities
proba, labels = torch.topk(pred_proba, k)
labels = labels.cpu().numpy().squeeze()
landmarks = [classes[i] for i in labels]
# Cleanup the labels
landmarks = [" ".join(label.split("_")) for label in landmarks]
proba = proba.squeeze().tolist()
fig = px.bar(
pd.DataFrame(data={"Landmarks": landmarks, "Probabilities": proba}),
y="Probabilities",
x="Landmarks",
)
return landmarks, proba, fig
|