🐶🧡🐱

app.py

@@ -0,0 +1,38 @@
+import torch
+import gradio as gr
+from app.model import PetClassificationModel
+from app.backbone import Backbone
+from app.config import CFG
+from torchvision import transforms
+
+# Load model
+backbone = Backbone(CFG.MODEL, len(CFG.idx_to_class), pretrained = CFG.PRETRAINED)
+model = PetClassificationModel(base_model = backbone.model, config = CFG)
+model.load_state_dict(torch.load('models/best_model.pt'))
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+# Eval mode
+model.eval()
+
+model.to(device)
+
+
+pred_transforms = transforms.Compose([
+    transforms.Resize(CFG.IMG_SIZE),
+    transforms.ToTensor(),
+])
+
+def predict(x):
+    x = pred_transforms(x).unsqueeze(0) # transform and batched
+    x = x.to(device)
+    
+    with torch.no_grad():
+        prediction = torch.nn.functional.softmax(model(x)[0], dim=0)
+        confidences = {CFG.idx_to_class[i]: float(prediction[i]) for i in range(len(CFG.idx_to_class))}
+        
+    return confidences
+
+gr.Interface(fn=predict,
+             inputs=gr.Image(type="pil"),
+             outputs=gr.Label(num_top_classes=5),
+             examples=["statics/pug.jpg", "statics/poodle.jpg", "statics/cat.jpg", "statics/no.jpg"]).launch()

app/backbone.py

@@ -0,0 +1,6 @@
+import timm
+from dataclasses import dataclass
+
+class Backbone:
+    def __init__(self, model, num_classes, pretrained = True):
+        self.model = timm.create_model(model, pretrained = pretrained, num_classes = num_classes)

app/config.py

@@ -0,0 +1,14 @@
+import pandas as pd
+
+class CFG:
+    LABEL_PATH = 'data/labels.csv'
+
+    labels = pd.read_csv(LABEL_PATH)
+    idx_to_class = dict(enumerate(labels.breed.unique()))
+    class_to_idx = {c:i for i,c in idx_to_class.items()}
+
+
+    # Model related
+    MODEL = 'inception_v4'
+    PRETRAINED = True
+    IMG_SIZE = (299, 299) # Depends in base model

app/model.py

@@ -0,0 +1,123 @@
+
+import torch
+from torch import nn
+
+
+import lightning as L
+
+import torch.nn.functional as F
+from torch import optim
+from torchmetrics import Accuracy
+
+from torch.optim.lr_scheduler import ReduceLROnPlateau
+
+
+
+class PetClassificationModel(L.LightningModule):
+  def __init__(self, base_model, config):
+    super().__init__()
+    self.config = config
+    self.num_classes = len(self.config.idx_to_class)
+    metric = Accuracy(task="multiclass", num_classes=self.num_classes)
+    self.train_acc = metric.clone()
+    self.val_acc = metric.clone()
+    self.test_acc = metric.clone()
+    self.training_step_outputs = []
+    self.validation_step_outputs = []
+    self.test_step_outputs = []
+    self.device_ = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+    self.pretrained_model = base_model
+    out_features = self.pretrained_model.get_classifier().out_features
+    self.custom_layers = nn.Sequential(
+          nn.Linear(out_features, 512, device = self.device_),
+          nn.ReLU(),
+          nn.Dropout(),
+          nn.Linear(512, self.num_classes, device = self.device_),
+        )
+
+  def forward(self, x):
+    x = self.pretrained_model(x)
+    #x = self.custom_layers(x)
+    return x
+
+
+  def training_step(self, batch, batch_idx):
+    x,y = batch
+    logits = self.forward(x) # -> logits
+    loss = F.cross_entropy(logits, y)
+    self.log_dict({'train_loss': loss})
+    self.training_step_outputs.append({'loss': loss, 'logits': logits, 'y':y})
+    return loss
+
+  def on_train_epoch_end(self):
+    # Concat batches
+    outputs = self.training_step_outputs
+    logits = torch.cat([x['logits'] for x in outputs])
+    y = torch.cat([x['y'] for x in outputs])
+    self.train_acc(logits, y)
+    self.log_dict({
+        'train_acc': self.train_acc,
+      },
+      on_step = False,
+      on_epoch = True,
+      prog_bar = True)
+    self.training_step_outputs.clear()
+
+  def validation_step(self, batch, batch_idx):
+    x,y = batch
+    logits = self.forward(x)
+    loss = F.cross_entropy(logits, y)
+    self.log_dict({'val_loss': loss})
+    self.validation_step_outputs.append({'loss': loss, 'logits': logits, 'y':y})
+    return loss
+
+  def on_validation_epoch_end(self):
+    # Concat batches
+    outputs = self.validation_step_outputs
+    logits = torch.cat([x['logits'] for x in outputs])
+    y = torch.cat([x['y'] for x in outputs])
+    self.val_acc(logits, y)
+    self.log_dict({
+        'val_acc': self.val_acc,
+      },
+      on_step = False,
+      on_epoch = True,
+      prog_bar = True)
+    self.validation_step_outputs.clear()
+
+  def test_step(self, batch, batch_idx):
+    x,y = batch
+    logits = self.forward(x)
+    loss = F.cross_entropy(logits, y)
+    self.log_dict({'test_loss': loss})
+    self.test_step_outputs.append({'loss': loss, 'logits': logits, 'y':y})
+    return loss
+
+  def on_test_epoch_end(self):
+    # Concat batches
+    outputs = self.test_step_outputs
+    logits = torch.cat([x['logits'] for x in outputs])
+    y = torch.cat([x['y'] for x in outputs])
+    self.test_acc(logits, y)
+    self.log_dict({
+        'test_acc': self.test_acc,
+      },
+      on_step = False,
+      on_epoch = True,
+      prog_bar = True)
+    self.test_step_outputs.clear()
+
+  def predict_step(self, batch):
+        x, y = batch
+        return self.model(x, y)
+
+  def configure_optimizers(self):
+    optimizer = optim.Adam(self.parameters(), lr=self.config.LEARNING_RATE)
+    lr_scheduler = ReduceLROnPlateau(optimizer, mode = 'min', patience = 3)
+    lr_scheduler_dict = {
+        "scheduler": lr_scheduler,
+        "interval": "epoch",
+         "monitor": "val_loss",
+    }
+    return {'optimizer': optimizer, 'lr_scheduler': lr_scheduler_dict}

model/best_model.pt

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

requirements.txt

@@ -0,0 +1,21 @@
+gradio==3.50.2
+gradio_client==0.6.1
+huggingface-hub==0.18.0
+lightning==2.1.0
+lightning-utilities==0.9.0
+mypy-extensions==1.0.0
+numpy==1.25.2
+pandas==2.1.1
+Pillow==10.1.0
+python-dateutil==2.8.2
+python-multipart==0.0.6
+pytorch-lightning==2.1.0
+rpds-py==0.10.6
+safetensors==0.4.0
+scikit-learn==1.3.1
+scipy==1.9.3
+timm==0.9.7
+torch==2.1.0
+torchmetrics==1.2.0
+torchvision==0.16.0
+