pushing app

.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
+*.jpeg filter=lfs diff=lfs merge=lfs -text
+*.jpg filter=lfs diff=lfs merge=lfs -text

README.md

@@ -5,7 +5,7 @@ colorFrom: yellow
 colorTo: blue
 sdk: gradio
 sdk_version: 3.19.1
-app_file: app.py
+app_file: app/app.py
 pinned: false
 license: mit
 ---

app/app.py

@@ -0,0 +1,101 @@
+import os
+import sys
+
+current = os.path.dirname(os.path.realpath(__file__))
+parent = os.path.dirname(current)
+sys.path.append(parent)
+
+import albumentations as A
+import gradio as gr
+import matplotlib.pyplot as plt
+import numpy as np
+import torch
+from albumentations.pytorch import ToTensorV2
+from PIL import Image
+
+from model import Classifier
+
+# Load the model
+model = Classifier.load_from_checkpoint("./models/checkpoint.ckpt")
+model.eval()
+
+# Define labels
+labels = [
+    "dog",
+    "horse",
+    "elephant",
+    "butterfly",
+    "chicken",
+    "cat",
+    "cow",
+    "sheep",
+    "spider",
+    "squirrel",
+]
+
+# Preprocess function
+def preprocess(image):
+    image = np.array(image)
+    resize = A.Resize(224, 224)
+    normalize = A.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5))
+    to_tensor = ToTensorV2()
+    transform = A.Compose([resize, normalize, to_tensor])
+    image = transform(image=image)["image"]
+    return image
+
+
+# Define the sample images
+sample_images = {
+    "dog": "./test_images/dog.jpeg",
+    "cat": "./test_images/cat.jpeg",
+    "butterfly": "./test_images/butterfly.jpeg",
+    "elephant": "./test_images/elephant.jpg",
+    "horse": "./test_images/horse.jpeg",
+}
+
+# Define the function to make predictions on an image
+def predict(image):
+    try:
+        image = preprocess(image).unsqueeze(0)
+
+        # Prediction
+        # Make a prediction on the image
+        with torch.no_grad():
+            output = model(image)
+            # convert to probabilities
+            probabilities = torch.nn.functional.softmax(torch.exp(output[0]), dim=0)
+            topk_prob, topk_label = torch.topk(probabilities, 3)
+
+            # Return the top 3 predictions
+        return {labels[i]: float(probabilities[i]) for i in range(3)}
+    except Exception as e:
+        print(f"Error predicting image: {e}")
+        return []
+
+
+# Define the interface
+def app():
+    title = "Animal-10 Image Classification"
+    description = "Classify images using a custom CNN model and deploy using Gradio."
+
+    gr.Interface(
+        title=title,
+        description=description,
+        fn=predict,
+        inputs=gr.Image(type="pil"),
+        outputs=gr.Label(
+            num_top_classes=3,
+        ),
+        examples=[
+            "./test_images/dog.jpeg",
+            "./test_images/cat.jpeg",
+            "./test_images/butterfly.jpeg",
+            "./test_images/elephant.jpg",
+            "./test_images/horse.jpeg",
+        ],
+    ).launch()
+
+
+# Run the app
+if __name__ == "__main__":
+    app()

config/args.py

@@ -0,0 +1,17 @@
+from dataclasses import dataclass
+
+
+@dataclass
+class Args:
+    """
+    Training arguments.
+    """
+
+    # Learning rate for the optimizer
+    learning_rate: float = 1e-3
+    # Training batch size
+    batch_size: int = 32
+    # Total numebr of classes
+    num_classes: int = 10
+    # Maximum number of training epochs
+    max_epochs: int = 5

model.py

@@ -0,0 +1,87 @@
+import pytorch_lightning as pl
+import torch
+import torchmetrics
+from simple_parsing import ArgumentParser
+from torch import nn
+from torch.nn import functional as F
+
+from config.args import Args
+
+parser = ArgumentParser()
+parser.add_arguments(Args, dest="options")
+args_namespace = parser.parse_args()
+args = args_namespace.options
+
+# Model class
+class Model(nn.Module):
+    def __init__(self):
+        super().__init__()
+
+        self.conv1 = nn.Conv2d(3, 32, 5)
+        self.conv2 = nn.Conv2d(32, 64, 5)
+        self.conv3 = nn.Conv2d(64, 128, 3)
+        self.dropout1 = nn.Dropout2d(0.25)
+        self.dropout2 = nn.Dropout2d(0.5)
+
+        x = torch.randn(3, 224, 224).view(-1, 3, 224, 224)
+        self._to_linear = None
+        self.convs(x)
+
+        self.fc1 = nn.Linear(self._to_linear, 128)
+        self.fc2 = nn.Linear(128, args.num_classes)
+
+    def convs(self, x):
+        x = F.max_pool2d(F.relu(self.conv1(x)), (2, 2))
+        x = self.dropout1(x)
+        x = F.max_pool2d(F.relu(self.conv2(x)), (2, 2))
+        x = self.dropout2(x)
+        x = F.max_pool2d(F.relu(self.conv3(x)), (2, 2))
+
+        if self._to_linear is None:
+            self._to_linear = x[0].shape[0] * x[0].shape[1] * x[0].shape[2]
+        return x
+
+    def forward(self, x):
+        x = self.convs(x)
+        x = x.view(-1, self._to_linear)
+        x = F.relu(self.fc1(x))
+        x = self.fc2(x)
+        return F.log_softmax(x, dim=1)
+
+
+class Classifier(pl.LightningModule):
+    def __init__(self):
+        super().__init__()
+
+        self.model = Model()
+        self.accuracy = torchmetrics.Accuracy(
+            task="multiclass", num_classes=args.num_classes
+        )
+
+    def forward(self, x):
+        x = self.model(x)
+        return x
+
+    def nll_loss(self, logits, labels):
+        return F.nll_loss(logits, labels)
+
+    def training_step(self, train_batch, batch_idx):
+        x, y = train_batch
+        logits = self.model(x)
+        loss = self.nll_loss(logits, y)
+        acc = self.accuracy(logits, y)
+        self.log("accuracy/train_accuracy", acc)
+        self.log("loss/train_loss", loss)
+        return loss
+
+    def validation_step(self, val_batch, batch_idx):
+        x, y = val_batch
+        logits = self.model(x)
+        loss = self.nll_loss(logits, y)
+        acc = self.accuracy(logits, y)
+        self.log("accuracy/val_accuracy", acc)
+        self.log("loss/val_loss", loss)
+
+    def configure_optimizers(self):
+        optimizer = torch.optim.Adam(self.parameters(), lr=args.learning_rate)
+        return optimizer

models/checkpoint.ckpt

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

requirements.txt

@@ -0,0 +1,6 @@
+torch
+torchvision
+pytorch-lightning
+simple-parsing
+albumentations
+matplotlib