import requests

import torch
import torchvision

import gradio as gr
import torch.nn.functional as F
import torchvision.transforms as T

# Get pre-trained model
weights = torchvision.models.ResNet18_Weights.IMAGENET1K_V1
model = torchvision.models.resnet18(weights=weights)

# Set model to evaluation mode
model.eval()

# Get ImageNet labels
r = requests.get("https://git.io/JJkYN")
labels = r.text.split("\n")

# Define prediction function
def predict(img):
    '''
    img: PIL image to be predicted
    confidences: python dictionary containing confidences
    '''
    # Transform image to pytorch tensor of shape [1, 3, 224, 224]
    # Put your image transforms here
    img = T.PILToTensor()(img).unsqueeze(0)
    img = T.Resize(size=(224, 224))(img)
    img = img / 255
    
    # Use model without gradients to reduce computation
    with torch.no_grad():
        # Get the tensor that contains the predictions
        out = model(img)
        prediction = F.softmax(out[0], dim=0)
        # Turn the above tensor into a dictionary with the humand-readable label and probability.
        confidences = {labels[i]: float(prediction[i]) for i in range(1000)}

    return confidences

gr.Interface(fn=predict,
             inputs=gr.Image(type="pil"),
             outputs=gr.Label(num_top_classes=10),
             theme="default",
            ).launch(share=False)