File size: 1,691 Bytes
a3c88bc a263b83 a3c88bc a263b83 a3c88bc a263b83 | 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 | import gradio as gr
import torch
import numpy as np
from model import *
def load_cub200_classes():
"""
This function loads the classes from the classes.txt file and returns a dictionary
"""
with open("classes.txt", encoding="utf-8") as f:
classes = f.read().splitlines()
# convert classes to dictionary separating the lines by the first space
classes = {int(line.split(" ")[0]) : line.split(" ")[1] for line in classes}
# return the classes dictionary
return classes
def load_model():
"""
This function loads the trained model and returns it
"""
# load the resnet model
model = resnet50(pretrained=False, stride=[1, 2, 2, 1], num_classes=200)
# load the trained weights
model.load_state_dict(torch.load("resnet.pt", map_location=torch.device('cpu')))
# set the model to evaluation mode
model.eval()
# return the model
return model
def predict_image(image):
"""
This function takes an image as input and returns the class label
"""
# load the model
model = load_model()
# load the classes
classes = load_cub200_classes()
# convert image to tensor
tensor = torch.from_numpy(image).permute(2, 0, 1).float().unsqueeze(0)
# make prediction
prediction = model(tensor).detach().numpy()[0]
# convert prediction to probabilities
probabilities = np.exp(prediction) / np.sum(np.exp(prediction))
# get the class with the highest probability
class_idx = np.argmax(probabilities)
# return the class label
return "Class: " + classes[class_idx]
# create a gradio interface
gr.Interface(fn=predict_image, inputs="image", outputs="text").launch()
|