Fixed Model

app.py

@@ -1,6 +1,4 @@
 import gradio as gr
-import os 
-import io
 import numpy as np
 import pandas as pd
 import matplotlib.pyplot as plt
@@ -8,7 +6,6 @@ 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
@@ -16,38 +13,11 @@ 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):
+class DogBreedPretrainedWideResnet(torch.nn.Module):
     def __init__(self):
         super().__init__()
         
@@ -62,60 +32,38 @@ class DogBreedPretrainedWideResnet(ImageClassificationBase):
     def forward(self, xb):
         return self.network(xb)
 
-def predict_single(img):
-    xb = transform_image(img) # adding extra dimension
-    xb = to_device(img, 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))
         
+            
+device = get_default_device()
+PATH = "./model/model.zip"
+map_location=torch.device('cpu')
+def predict_single(img):
+    xb = transform_image(img)           # Transforming image to Tensor
+    xb = xb.to(device)
+    preds = model(xb)                   # change model object here
+    max_val, kls = torch.max(preds, 1)
+    print('Predicted :', breeds[kls])
+    return breeds[kls]
+
 def image_mod(image):
-    im_resize = image.resize((500, 500))
-    buf = io.BytesIO()
-    im_resize.save(buf, format='JPEG')
-    byte_im = buf.getvalue()
-    return predict_single(byte_im)
+    return predict_single(image)
     
 def transform_image(image_bytes):
-    my_transforms = transforms.Compose([transforms.Resize(255),
-                                        transforms.CenterCrop(224),
+    my_transforms = transforms.Compose([transforms.Resize((500)),
                                         transforms.ToTensor(),
                                         transforms.Normalize(
                                             [0.485, 0.456, 0.406],
                                             [0.229, 0.224, 0.225])])
-    image = Image.open(io.BytesIO(image_bytes))
-    return my_transforms(image).unsqueeze(0)
-    
-device = get_default_device()
-PATH = "./model/model.zip"
-map_location=torch.device('cpu')
-test_transform = transforms.Compose([
-    transforms.Resize((224,224)), 
-    transforms.ToTensor(),
-#    transforms.Normalize(*imagenet_stats, inplace=True)
-])
+    return my_transforms(image_bytes).unsqueeze(0)
+
 model = DogBreedPretrainedWideResnet()
 model.load_state_dict(torch.load(PATH,map_location))
+model.eval()
 breeds=['Chihuahua',
  'Japanese spaniel',
  'Maltese dog',
@@ -236,6 +184,6 @@ breeds=['Chihuahua',
  'dingo',
  'dhole',
  'African hunting dog']
-iface = gr.Interface(image_mod, gr.Image(type="pil"), "image")
+iface = gr.Interface(image_mod, gr.Image(type="pil"), "text", examples=["doggo1.jpg","doggo2.jpg","doggo3.png","doggo4.png"])
 
 iface.launch()

model/model.zip

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:25145664111db820344b774c5f1b693e8daeaf119dafc5a55d2890c83a438126
-size 268701569
+oid sha256:c7b67545692a07c4d4e255343f7fa5ff6a327c56ac4bcca38efc2867c8eb8a38
+size 268684285