BillyCoder13/Multi_Class_Image_Classification

1import gradio as gr2import torch3import numpy as np4 5from model import *6 7def load_cub200_classes():8    """9    This function loads the classes from the classes.txt file and returns a dictionary10    """11    with open("classes.txt", encoding="utf-8") as f:12        classes = f.read().splitlines()13 14    # convert classes to dictionary separating the lines by the first space15    classes = {int(line.split(" ")[0]) : line.split(" ")[1] for line in classes}16 17    # return the classes dictionary18    return classes19 20def load_model():21    """22    This function loads the trained model and returns it23    """24 25    # load the resnet model26    model = resnet50(pretrained=False, stride=[1, 2, 2, 1], num_classes=200)27    # load the trained weights28    model.load_state_dict(torch.load("resnet.pt", map_location=torch.device('cpu')))29    # set the model to evaluation mode30    model.eval()31    # return the model32    return model33 34def predict_image(image):35    """36    This function takes an image as input and returns the class label37    """38 39    # load the model40    model = load_model()41    # load the classes42    classes = load_cub200_classes()43 44    # convert image to tensor45    tensor = torch.from_numpy(image).permute(2, 0, 1).float().unsqueeze(0)46    # make prediction47    prediction = model(tensor).detach().numpy()[0]48    # convert prediction to probabilities49    probabilities = np.exp(prediction) / np.sum(np.exp(prediction))50    # get the class with the highest probability51    class_idx = np.argmax(probabilities)52    # return the class label53    return "Class: " + classes[class_idx]54 55# create a gradio interface56gr.Interface(fn=predict_image, inputs="image", outputs="text").launch()57