uploading app and initial commit

.gitattributes

@@ -32,3 +32,5 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+examples/Image_4.jpg filter=lfs diff=lfs merge=lfs -text
+examples/Image_89.jpg filter=lfs diff=lfs merge=lfs -text

app.py

@@ -0,0 +1,118 @@
+### 1. Imports and class names setup ###
+import gradio as gr
+import os
+import torch
+from torchvision import transforms
+
+from models import get_mobilenet_v2_model, get_resnet_18_model, get_vgg_16_model
+from timeit import default_timer as timer
+from typing import Tuple, Dict
+
+# Set device
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+# Setup class names
+class_names = ["car","dragon","hourse","pegasus","ship","t-rex","tree"]
+
+### 2. Model and transforms preparation ###
+
+# Create EffNetB2 model
+img_transforms = transforms.Compose(
+    [
+        transforms.Resize(size=(224, 224)),
+        transforms.ToTensor(),
+    ]
+)
+
+model_name_to_fn = {
+    "mobilenet_v2": get_mobilenet_v2_model,
+    "resnet_18": get_resnet_18_model,
+    "vgg_16": get_vgg_16_model,
+}
+model_name_to_path = {
+    "mobilenet_v2": "mobilenet_v2.pth",
+    "resnet_18": "resnet_18.pth",
+    "vgg_16": "vgg_16.pt",
+}
+
+### 3. Predict function ###
+
+
+# Create predict function
+def predict(img, model_name: str,) -> Tuple[Dict, float]:
+    """
+    Desc: Transforms and performs a prediction on img and returns prediction and time taken.
+    Args:
+        model_name (str): Name of the model to use for prediction.
+        img (PIL.Image): Image to perform prediction on.
+    Returns:
+        Tuple[Dict, float]: Tuple containing a dictionary of prediction labels and probabilities and the time taken to perform the prediction.
+    """
+    # Start the timer
+    start_time = timer()
+
+    # Get the model function based on the model name
+    model_fn = model_name_to_fn[model_name]
+    model_path = model_name_to_path[model_name]
+
+    # Create the model and load its weights
+    model = model_fn().to(device)
+    model.load_state_dict(
+        torch.load(f"./models/{model_name}.pth", map_location=torch.device(device=device))
+    )
+
+    # Put model into evaluation mode and turn on inference mode
+    model.eval()
+    with torch.inference_mode():
+        # Transform the target image and add a batch dimension
+        img = img_transforms(img).unsqueeze(0).to(device)
+
+        # Pass the transformed image through the model and turn the prediction logits into prediction probabilities
+        pred_probs = torch.softmax(model(img), dim=1)
+
+    # Create a prediction label and prediction probability dictionary for each prediction class (this is the required format for Gradio's output parameter)
+    pred_labels_and_probs = {
+        class_names[i]: float(pred_probs[0][i]) for i in range(len(class_names))
+    }
+
+    # Calculate the prediction time
+    pred_time = round(timer() - start_time, 5)
+
+    # Return the prediction dictionary and prediction time
+    return pred_labels_and_probs, pred_time
+
+
+### 4. Gradio app ###
+
+# Create title, description and article strings
+title = "SketchRec Mini ✍🏻"
+description = "An Mutimodel Sketch Recognition App 🎨"
+article = ""
+
+# Create examples list from "examples/" directory
+example_list = [["examples/" + example] for example in os.listdir("examples")]
+
+# Create the Gradio demo
+model_selection_dropdown = gr.components.Dropdown(
+    choices=list(model_name_to_fn.keys()), label="Select a model",
+    value="mobilenet_v2"
+)
+
+demo = gr.Interface(
+    fn=predict,  # mapping function from input to output
+    inputs=[gr.Image(type="pil"),model_selection_dropdown],  # what are the inputs?
+    outputs=[
+        gr.Label(num_top_classes=7, label="Predictions"),  # what are the outputs?
+        gr.Number(label="Prediction time (s)"),
+    ],  # our fn has two outputs, therefore we have two outputs
+    # Create examples list from "examples/" directory
+    examples=example_list,
+    title=title,
+    description=description,
+    article=article,
+)
+
+# Launch the demo!
+demo.launch(
+    debug=True,
+)

models.py

@@ -0,0 +1,183 @@
+from typing import List
+import torch
+from torch import nn
+import numpy as np
+from torchvision import models
+from torchvision.models import ResNet18_Weights,ResNet50_Weights,VGG16_Weights,MobileNet_V2_Weights
+
+
+class EarlyStopping:
+    def __init__(self, tolerance=5, min_delta=0):
+
+        self.tolerance = tolerance
+        self.min_delta = min_delta
+        self.counter = 0
+        self.early_stop = False
+
+    def __call__(self, train_loss, validation_loss):
+        if (validation_loss - train_loss) > self.min_delta: 
+            self.counter +=1
+            if self.counter >= self.tolerance:  
+                self.early_stop = True
+                
+                
+
+class Resnet18(nn.Module):
+    def __init__(self,out_shape:int = 1000) -> None:
+        super().__init__()
+        self.resnet = models.resnet18()
+        self.resnet.fc = nn.Linear(512,out_shape)
+    
+    def forward(self,x):
+        return self.resnet(x)
+    
+    
+class PretrainedResnet18(nn.Module):
+    def __init__(self,out_shape:int = 1000) -> None:
+        super().__init__()
+        self.resnet = models.resnet18(weights=ResNet18_Weights.DEFAULT)
+        
+        # freeze all layers except last fc layer
+        for parms in self.resnet.parameters():
+            parms.requires_grad = False
+        
+        self.resnet.fc = nn.Linear(512,out_shape)
+    
+    def forward(self,x):
+        return self.resnet(x)
+    
+
+class Resnet50(nn.Module):
+    def __init__(self,out_shape:int = 1000) -> None:
+        super().__init__()
+        self.resnet = models.resnet50()
+        self.resnet.fc = nn.Linear(2048,out_shape)
+    
+    def forward(self,x):
+        return self.resnet(x)
+    
+    
+class PretrainedResnet50(nn.Module):
+    def __init__(self,out_shape:int = 1000) -> None:
+        super().__init__()
+        self.resnet = models.resnet50(weights=ResNet50_Weights.DEFAULT)
+        
+        # freeze all layers except last fc layer
+        for parms in self.resnet.parameters():
+            parms.requires_grad = False
+        
+        self.resnet.fc = nn.Linear(2048,out_shape)
+    
+    def forward(self,x):
+        return self.resnet(x)
+    
+
+class EfficentNetB0(nn.Module):
+    def __init__(self,out_shape:int = 1000) -> None:
+        super().__init__()
+        self.effnet = models.efficientnet_b0()
+        self.effnet.classifier = nn.Linear(1280,out_shape)
+    
+    def forward(self,x):
+        return self.effnet(x)
+
+
+class MobileNetV2(nn.Module):
+    def __init__(self,out_shape:int = 1000) -> None:
+        super().__init__()
+        self.mobilenet = models.mobilenet_v2()
+        self.mobilenet.classifier[1] = nn.Linear(1280,out_shape)
+    
+    def forward(self,x):
+        return self.mobilenet(x)
+
+class PretrainedMobileNetV2(nn.Module):
+    def __init__(self,out_shape:int = 1000) -> None:
+        super().__init__()
+        self.mobilenet = models.mobilenet_v2(weights=MobileNet_V2_Weights.DEFAULT)
+        
+        # freeze all layers except last fc layer
+        for parms in self.mobilenet.parameters():
+            parms.requires_grad = False
+        
+        self.mobilenet.classifier = nn.Sequential(
+            nn.Dropout(p=0.2, inplace=False),
+            nn.Linear(in_features=1280, out_features=1000, bias=True)
+        )
+    
+    def forward(self,x):
+        return self.mobilenet(x)
+
+
+class VGG16(nn.Module):
+    def __init__(self,out_shape:int = 1000) -> None:
+        super().__init__()
+        self.vgg = models.vgg16()
+        self.vgg.classifier = nn.Sequential(
+            nn.Linear(in_features=25088, out_features=4096, bias=True),
+            nn.ReLU(inplace=True),
+            nn.Dropout(p=0.5, inplace=False),
+            nn.Linear(in_features=4096, out_features=4096, bias=True),
+            nn.ReLU(inplace=True),
+            nn.Dropout(p=0.5, inplace=False),
+            nn.Linear(in_features=4096, out_features=out_shape, bias=True),
+        )
+    
+    def forward(self,x):
+        return self.vgg(x)
+    
+    
+class PretrainedVGG16(nn.Module):
+    def __init__(self,out_shape:int = 1000) -> None:
+        super().__init__()
+        self.vgg = models.vgg16(weights=VGG16_Weights.DEFAULT)
+        
+        # freeze all layers except last clf layer
+        for parms in self.vgg.parameters():
+            parms.requires_grad = False
+        
+        self.vgg.classifier = nn.Sequential(
+            nn.Linear(in_features=25088, out_features=4096, bias=True),
+            nn.ReLU(inplace=True),
+            nn.Dropout(p=0.5, inplace=False),
+            nn.Linear(in_features=4096, out_features=4096, bias=True),
+            nn.ReLU(inplace=True),
+            nn.Dropout(p=0.5, inplace=False),
+            nn.Linear(in_features=4096, out_features=out_shape, bias=True),
+        )
+    
+    def forward(self,x):
+        return self.vgg(x)
+    
+
+
+class VIT(nn.Module):
+    def __init__(self,out_shape:int = 1000) -> None:
+        super().__init__()
+        self.vit = models.vit_b_16()
+        self.vit.head = nn.Linear(768,out_shape)
+    
+    def forward(self,x):
+        return self.vit(x)
+
+
+
+# functions to get models
+def get_resnet_18_model():
+    model = Resnet18(out_shape=7)
+    return model
+
+
+def get_resnet_50_model():
+    model = Resnet50(out_shape=7)
+    return model
+
+
+def get_vgg_16_model():
+    model = VGG16(out_shape=7)
+    return model
+
+
+def get_mobilenet_v2_model():
+    model = MobileNetV2(out_shape=7)
+    return model

models/mobilenet_v2.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7597e7b6ba1183b33aca66ebf39eb73359bbc568b7293e2e62bc52cdad3f48e7
+size 9263285

models/resnet_18.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ce03f8579e35fd33405cf5ac1b30fe62f87a8b47dc8c635fdaa500398bd915b2
+size 44822093

models/vgg_16.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4c5f6a03dbaaa82674d76eaf0fd7fb831633d2929a25e15040ca3f05d4761f1b
+size 537176517

requirements.txt

@@ -0,0 +1,3 @@
+torch
+torchvision
+gradio