Spaces:
Build error
Build error
import gradio as gr | |
import os | |
import numpy as np | |
import pandas as pd | |
import matplotlib.pyplot as plt | |
from PIL import Image | |
from sklearn.preprocessing import LabelEncoder | |
import torch | |
import torch.nn.functional as F | |
import torchvision | |
from torchvision import transforms | |
import torchvision.models as models | |
from torchvision.datasets import ImageFolder | |
from torch.utils.data.dataset import Dataset | |
from torch.utils.data import Dataset, random_split, DataLoader | |
from torch.utils.data import DataLoader | |
from sklearn.model_selection import train_test_split | |
import torchmetrics | |
from tqdm.notebook import tqdm | |
class ImageClassificationBase(torch.nn.Module): | |
# training step | |
def training_step(self, batch): | |
img, targets = batch | |
out = self(img) | |
loss = F.nll_loss(out, targets) | |
return loss | |
# validation step | |
def validation_step(self, batch): | |
img, targets = batch | |
out = self(img) | |
loss = F.nll_loss(out, targets) | |
acc = accuracy(out, targets) | |
return {'val_acc':acc.detach(), 'val_loss':loss.detach()} | |
# validation epoch end | |
def validation_epoch_end(self, outputs): | |
batch_losses = [x['val_loss'] for x in outputs] | |
epoch_loss = torch.stack(batch_losses).mean() | |
batch_accs = [x['val_acc'] for x in outputs] | |
epoch_acc = torch.stack(batch_accs).mean() | |
return {'val_loss':epoch_loss.item(), 'val_acc':epoch_acc.item()} | |
# print result end epoch | |
def epoch_end(self, epoch, result): | |
print("Epoch [{}] : train_loss: {:.4f}, val_loss: {:.4f}, val_acc: {:.4f}".format(epoch, result["train_loss"], result["val_loss"], result["val_acc"])) | |
class DogBreedPretrainedWideResnet(ImageClassificationBase): | |
def __init__(self): | |
super().__init__() | |
self.network = models.wide_resnet50_2(pretrained=True) | |
# Replace last layer | |
num_ftrs = self.network.fc.in_features | |
self.network.fc = torch.nn.Sequential( | |
torch.nn.Linear(num_ftrs, 120), | |
torch.nn.LogSoftmax(dim=1) | |
) | |
def forward(self, xb): | |
return self.network(xb) | |
def predict_single(img): | |
xb = img.unsqueeze(0) # adding extra dimension | |
xb = to_device(xb, device) | |
preds = model(xb) # change model object here | |
predictions = preds[0] | |
max_val, kls = torch.max(predictions, dim=0) | |
print('Predicted :', breeds[kls]) | |
plt.imshow(img.permute(1,2,0)) | |
plt.show() | |
def get_default_device(): | |
if torch.cuda.is_available(): | |
return torch.device('cuda') | |
else: | |
return torch.device('cpu') | |
def to_device(data, device): | |
if isinstance(data, (list, tuple)): | |
return [to_device(d, device) for d in data] | |
else: | |
return data.to(device, non_blocking=True) | |
def accuracy(outputs, labels): | |
_, preds = torch.max(outputs, dim=1) | |
return torch.tensor(torch.sum(preds == labels).item() / len(preds)) | |
def image_mod(image): | |
return predict_single(image) | |
device = get_default_device() | |
PATH = "./model/model.zip" | |
model = DogBreedPretrainedWideResnet() | |
model.load_state_dict(torch.load(PATH)) | |
breeds=['Chihuahua', | |
'Japanese spaniel', | |
'Maltese dog', | |
'Pekinese', | |
'Shih Tzu', | |
'Blenheim spaniel', | |
'papillon', | |
'toy terrier', | |
'Rhodesian ridgeback', | |
'Afghan hound', | |
'basset', | |
'beagle', | |
'bloodhound', | |
'bluetick', | |
'black and tan coonhound', | |
'Walker hound', | |
'English foxhound', | |
'redbone', | |
'borzoi', | |
'Irish wolfhound', | |
'Italian greyhound', | |
'whippet', | |
'Ibizan hound', | |
'Norwegian elkhound', | |
'otterhound', | |
'Saluki', | |
'Scottish deerhound', | |
'Weimaraner', | |
'Staffordshire bullterrier', | |
'American Staffordshire terrier', | |
'Bedlington terrier', | |
'Border terrier', | |
'Kerry blue terrier', | |
'Irish terrier', | |
'Norfolk terrier', | |
'Norwich terrier', | |
'Yorkshire terrier', | |
'wire haired fox terrier', | |
'Lakeland terrier', | |
'Sealyham terrier', | |
'Airedale', | |
'cairn', | |
'Australian terrier', | |
'Dandie Dinmont', | |
'Boston bull', | |
'miniature schnauzer', | |
'giant schnauzer', | |
'standard schnauzer', | |
'Scotch terrier', | |
'Tibetan terrier', | |
'silky terrier', | |
'soft coated wheaten terrier', | |
'West Highland white terrier', | |
'Lhasa', | |
'flat coated retriever', | |
'curly coated retriever', | |
'golden retriever', | |
'Labrador retriever', | |
'Chesapeake Bay retriever', | |
'German short haired pointer', | |
'vizsla', | |
'English setter', | |
'Irish setter', | |
'Gordon setter', | |
'Brittany spaniel', | |
'clumber', | |
'English springer', | |
'Welsh springer spaniel', | |
'cocker spaniel', | |
'Sussex spaniel', | |
'Irish water spaniel', | |
'kuvasz', | |
'schipperke', | |
'groenendael', | |
'malinois', | |
'briard', | |
'kelpie', | |
'komondor', | |
'Old English sheepdog', | |
'Shetland sheepdog', | |
'collie', | |
'Border collie', | |
'Bouvier des Flandres', | |
'Rottweiler', | |
'German shepherd', | |
'Doberman', | |
'miniature pinscher', | |
'Greater Swiss Mountain dog', | |
'Bernese mountain dog', | |
'Appenzeller', | |
'EntleBucher', | |
'boxer', | |
'bull mastiff', | |
'Tibetan mastiff', | |
'French bulldog', | |
'Great Dane', | |
'Saint Bernard', | |
'Eskimo dog', | |
'malamute', | |
'Siberian husky', | |
'affenpinscher', | |
'basenji', | |
'pug', | |
'Leonberg', | |
'Newfoundland', | |
'Great Pyrenees', | |
'Samoyed', | |
'Pomeranian', | |
'chow', | |
'keeshond', | |
'Brabancon griffon', | |
'Pembroke', | |
'Cardigan', | |
'toy poodle', | |
'miniature poodle', | |
'standard poodle', | |
'Mexican hairless', | |
'dingo', | |
'dhole', | |
'African hunting dog'] | |
iface = gr.Interface(image_mod, gr.Image(type="pil"), "image") | |
iface.launch() |