import gradio as gr
import numpy as np
import torch

learn_inf = torch.jit.load("checkpoints/transfer_exported.pt")

def classify(img):
    
    # Transform to tensor
    timg = torch.from_numpy(img.transpose((2,0,1))).unsqueeze_(0)
    
    # Calling the model
    softmax = learn_inf(timg).data.cpu().numpy().squeeze()
    
    # Get the indexes of the classes ordered by softmax
    # (larger first)
    idxs = np.argsort(softmax)[::-1]
    
    # Loop over the classes with the largest softmax
    label = ""
    for i in range(5):
        # Get softmax value
        p = softmax[idxs[i]]
        
        # Get class name
        landmark_name = learn_inf.class_names[idxs[i]]
        
        label += f"{landmark_name[3:]} (prob: {p:.2f})" + "\n"
    return label

app = gr.Interface(
    fn=classify,
    inputs=gr.Image(),
    outputs=["text"],
    api_name="classify",
    title="Landmark Classifier",
    description="Upload an image of a landmark and this program will guess what landmark is in the picture.",
    # examples="static_images/examples"
)

app.launch()