File size: 940 Bytes
6c1aaee |
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 |
# utils.py
import torch
import json
from torchvision import transforms
with open('label_mapping.json', 'r') as json_file:
label_mapping = json.load(json_file)
def load_model(path):
model = torch.jit.load(path, map_location=torch.device("cpu"))
return model
def predict(model, image):
model.eval()
# Transform the image
transform = transforms.Compose([transforms.Resize((224, 224)), transforms.ToTensor()])
image = transform(image)
with torch.no_grad():
image = image.unsqueeze(0)
output = model(image)
probabilities = torch.nn.functional.softmax(output, dim=1)
_, predicted_class = torch.max(probabilities, 1)
# Convert predicted class index to label name using label_mapping
predicted_label = label_mapping[f"{predicted_class.item()}"]
probability= probabilities[0][predicted_class].item()
return predicted_label, round(probability, 2) |