Spaces:
Build error
Build error
import requests | |
import gradio as gr | |
import torch | |
from timm.data import resolve_data_config | |
from timm.data.transforms_factory import create_transform | |
LABELS = {0:'Cat', 1:'Dog'} | |
model = torch.load('CatVsDogsModel.pth',map_location='cpu') | |
transform = create_transform(**resolve_data_config({},model=model)) | |
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=2) | |
return {LABELS[i]: v.item() for i,v in zip(indices,values)} | |
transform = create_transform(**resolve_data_config({},model=model)) | |
# we do not need to train model , hence using model.eval() to use it only for inference | |
model.eval() | |
iface = gr.Interface(fn=predict, inputs=gr.inputs.Image(type='pil'), outputs="label").launch() | |
iface.launch() | |