import torch
from models.common import DetectMultiBackend
from torchvision import transforms
import gradio as gr
import requests
from PIL import Image

weights='/content/drive/MyDrive/yolov5/yolov5s-cls.pt'

model = DetectMultiBackend(weights)

# load imagenet 1000 labels
response = requests.get("https://git.io/JJkYN")
labels = response.text.split("\n")

def preprocess_image(inp):
    # Define the preprocessing steps
    preprocess = transforms.Compose([
        transforms.Resize(256),
        transforms.CenterCrop(224),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
    ])
    # Apply the preprocessing steps to the image
    image = preprocess(inp)
    # Convert the image to a PyTorch tensor
    image = torch.tensor(image).unsqueeze(0)

    return image

def predict(inp):

  with torch.no_grad():
    prediction = torch.nn.functional.softmax(model(preprocess_image(inp))[0], dim=0)

    print(prediction)
    confidences = {labels[i]: float(prediction[i]) for i in range(1000)}    
  return confidences


gr.Interface(fn=predict, 
             inputs=gr.Image(type="pil"),
             outputs="label",labels=labels).launch(debug=True)
            
             #outputs=gr.Label(num_top_classes=5))