Sync from Github Actions

app.py

@@ -1,57 +1,21 @@
 import gradio as gr
-import numpy as np
-import torch
-import torchvision.transforms.functional as TF
-from model.cnn import CNN
+from model.predict import predict
+from torchvision.transforms import functional as F
 
-path = "model/epoch=599-step=187800.ckpt"
-model = CNN.load_from_checkpoint(path)
-model.eval()
+def custom_predict(image):
+    image = F.to_tensor(image)
+    return predict(image, get_dictionary=True)
 
 
-translation = {
-    "airplane": 0,
-    "automobile": 1,
-    "bird": 2,
-    "cat": 3,
-    "deer": 4,
-    "dog": 5,
-    "frog": 6,
-    "horse": 7,
-    "ship": 8,
-    "truck": 9,
-}
-inverted_translation = {v: k for k, v in translation.items()}
-
-
-def predict(image):
-    
-    image = TF.to_tensor(image)
-    image_tensor = torch.stack([image])  
-
-    print(image_tensor) 
-    
-    result = model(image_tensor)
-    result = torch.softmax(result,dim=1)
-    result = result[0]
-
-
-    dict_results = {}
-    
-    for i in range(len(result)):
-        dict_results[inverted_translation[i]] = float(result[i])
-
-    return dict_results
-
 
 demo = gr.Interface(
-    predict,
+    custom_predict,
     title="Image Classifier using CNN ( Cifar-10) ",
     description="This is a image classifier using a CNN, it was trained on the Cifar-10 dataset ( Kaggle) \n",
-    article="The architecture is a CNN",
+    article="The architecture is a CNN, uploaded via Github Actions",
     inputs=gr.Image(shape=(32, 32),type="pil"),
     outputs=gr.Label(),
-    examples=["examples/1.png", "examples/2.png", "examples/3.png", "examples/4.png" ],
+    examples=["examples/1.png", "examples/2.png", "examples/3.png", "examples/4.png" , "examples/5.png"],
 )
 
-demo.launch()
+demo.launch()

dataset/__init__.py

@@ -0,0 +1 @@
+from dataset.dataset import CifarDataset

dataset/dataset.py

@@ -0,0 +1,31 @@
+import pandas as pd
+from torch.utils.data import Dataset
+from torchvision.io import read_image
+
+translation = {'airplane': 0, 'automobile': 1, 'bird': 2, 'cat': 3, 'deer': 4, 'dog': 5, 'frog': 6, 'horse': 7, 'ship': 8, 'truck': 9}
+inverted_translation = {v: k for k, v in translation.items()}
+
+
+class CifarDataset(Dataset):
+    def __init__(self, file, folder, include_idx=False):
+        self.data = pd.read_csv(file)
+        self.data['label'] = self.data['label'].map(translation)
+        self.folder = folder
+        self.include_idx = include_idx
+
+    def __len__(self):
+        return len(self.data)
+    
+    def __getitem__(self, idx):
+        id, label = self.data.iloc[idx][0] , self.data.iloc[idx][1]
+        image = read_image( f'{self.folder}/{id}.png')
+        image = image/255.0
+        
+        if self.include_idx:
+            return image, label, id
+        else:
+            return image, label
+
+if __name__ == '__main__':
+    dataset = CifarDataset('trainLabels.csv', 'train')
+    print(dataset[0][0].type())

model/cnn.py

@@ -3,7 +3,6 @@ import torch
 import torch.nn as nn
 
 
-
 class CNN(pl.LightningModule):    
     def __init__(self):
         super().__init__()

model/predict.py

@@ -0,0 +1,30 @@
+from model.cnn import CNN
+import torch
+from dataset.dataset import inverted_translation
+import matplotlib.pyplot as plt 
+from torchvision.transforms import functional as F
+
+path = "model/epoch=599-step=187800.ckpt"
+model = CNN.load_from_checkpoint(path)
+model.eval()
+
+def predict(image, get_dictionary=False):
+    
+    image_tensor = image.view(1, 3, 32, 32)
+    result = model(image_tensor)
+    result = torch.softmax(result,dim=1)
+    result = result[0]
+
+    if get_dictionary:
+        dict_results = {}
+        
+        for i in range(len(result)):
+            dict_results[inverted_translation[i]] = float(result[i])
+
+        return dict_results
+
+    else:
+        best = int(torch.argmax(result))        
+        return inverted_translation[best]
+
+