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