import torch

model = torch.hub.load('pytorch/vision:v0.6.0', 'resnet18', pretrained=True).eval()


import requests
from PIL import Image
from torchvision import transforms

# Download human-readable labels for ImageNet.
response = requests.get("https://git.io/JJkYN")
labels = response.text.split("\n")

def predict(inp):
  inp = transforms.ToTensor()(inp).unsqueeze(0)
  with torch.no_grad():
    prediction = torch.nn.functional.softmax(model(inp)[0], dim=0)
    confidences = {labels[i]: float(prediction[i]) for i in range(1000)}
  return confidences


import gradio as gr

with gr.Blocks(title="Image Classification for 1000 Objects", css=".gradio-container {background:mintcream;}") as demo:
    gr.HTML("""<div style="font-family:'Times New Roman', 'Serif'; font-size:16pt; font-weight:bold; text-align:center; color:royalblue;">Image Classification for 1000 Objects</div>""")
    
    with gr.Row(): 
        input_image  = gr.Image(type="filepath", image_mode="L")        
        output_label = gr.Label(label="Probabilities", num_top_classes=3)
    
    send_btn = gr.Button("Infer")
    send_btn.click(fn=predict, inputs=input_image, outputs=output_label)
    
    with gr.Row():
        gr.Examples(['./lion.jpg']   , label='Sample images : Lion', inputs=input_image)
        gr.Examples(['./cheetah.jpg'], label='Cheetah'             , inputs=input_image)

# gr.Interface(fn=predict,
#              inputs=gr.Image(type="pil"),
#              outputs=gr.Label(num_top_classes=3),
#              examples=["lion.jpg", "cheetah.jpg"]).launch()

demo.launch(debug=True, share=True)