import torch
from PIL import Image
import torchvision.transforms as transforms
import numpy as np
import json
import requests
import matplotlib.pyplot as plt

import gradio as gr

device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")

efficientnet = torch.hub.load('NVIDIA/DeepLearningExamples:torchhub', 'nvidia_efficientnet_b0', pretrained=True)
utils = torch.hub.load('NVIDIA/DeepLearningExamples:torchhub', 'nvidia_convnets_processing_utils')

efficientnet.eval().to(device)

def inference(img):
  batch = torch.cat(
    [utils.prepare_input_from_uri(img)]
  ).to(device)
  with torch.no_grad():
      output = torch.nn.functional.softmax(efficientnet(batch), dim=1)
      
    
  results = utils.pick_n_best(predictions=output, n=5)
  
  return results
  
gr.Interface(inference,gr.inputs.Image(type="file"),"text").launch()