Spaces:
Sleeping
Sleeping
import requests | |
import gradio as gr | |
import torch | |
from timm import create_model | |
from timm import resolve_data_config | |
from timm.data.transforms_factory import create_transform | |
IMAGENET_1k_URL = "https://storage.googleapis.com/bit_models/ilsvrc2012_wordnet_lemmas.txt" | |
LABELS = requests.get() | |
model = create_model('resnet50',pretrained=True) | |
transform = create_transform(**resolve_data_config({},model=model)) | |
model.eval() | |
def predict(img): | |
img = img.convert('RGB') | |
img = transform(img).unsqueeze(0) | |
with torch.no_grad(): | |
out= model(img) | |
probability = torch.nn.functional.softmax(out[0],dim=0) | |
values, indices = torch.topk(probability,k=5) | |
return {LABELS[i]: v.items() for i,v in zip(indices,values)} | |
iface = gr.Interface(fn=predict, inputs=gr.inputs.Image(type='pil'), outputs="label").launch() | |
iface.launch() | |