Spaces:
Build error
Build error
| import requests | |
| import gradio as gr | |
| import torch | |
| from torchvision.models import alexnet | |
| import torch.nn as nn | |
| from torchvision import transforms | |
| LABELS = {0:'Cat', 1:'Dog'} | |
| model = alexnet(pretrained=True) | |
| for param in model.parameters(): | |
| param.requires_grad = False | |
| # # Add a avgpool here | |
| # avgpool = nn.AdaptiveAvgPool2d((7, 7)) | |
| # | |
| # # Replace the classifier layer | |
| # # to customise it according to our output | |
| model.classifier = nn.Sequential( | |
| nn.Linear( 256*7*7, 1024), | |
| nn.Linear(1024, 256), | |
| nn.Linear(256, 2)) | |
| checkpoint = torch.load( | |
| "CatVsDogsModel.pth", map_location=torch.device("cpu") | |
| ) | |
| model.load_state_dict(checkpoint["state_dict"]) | |
| model = model.to('cpu') | |
| transform = transforms.Compose( | |
| [transforms.Resize((128, 128)), transforms.ToTensor()] | |
| ) | |
| def predict(img): | |
| img = transform(img).to('cpu') | |
| img = img.unsqueeze(0) | |
| with torch.no_grad(): | |
| out= model(img) | |
| probability = torch.nn.functional.softmax(out[0],dim=0) | |
| print(out) | |
| print(type(img)) | |
| values, indices = torch.topk(probability,k=2) | |
| return {LABELS[i]: v.item() for i,v in zip(indices,values)} | |
| iface = gr.Interface(fn=predict, inputs=gr.inputs.Image(type='pil'), outputs="label").launch() | |
| iface.launch(share=True) | |