EfficientNet / app.py
akhaliq's picture
akhaliq HF Staff
Create app.py
679a7db
raw
history blame
859 Bytes
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()