[feat]: update model weight and deployment

.gitignore

@@ -209,6 +209,7 @@ __marimo__/
 
 # local
 datasets/*
-checkpoints/*
+checkpoints/ckpt_23_10_2025_1/*
 dc_env/*
 note.md 
+

deployment/gradio/main.py

@@ -1,30 +1,39 @@
 import gradio as gr
-
-
-def greet(name):
-    return f"Hello {name}!"
-
-def classify_image(image):
-    return "cat"  # Placeholder for actual image classification logic
-
+from src.infer import inference_pipeline
+
+model_path = "checkpoints/ckpt_23_10_2025/best_cat_dog_classifier_model_20251019_122336.pth"
+
+def classify_image(
+        image_path: str
+) -> str:
+    """
+    Classify the input image as cat or dog.
+    """
+    if image_path is None:
+        return "Please upload an image."
+    try:
+        prediction = inference_pipeline(
+            image_path=image_path, 
+            model_path=model_path
+        )
+        return f"Prediction: {prediction.capitalize()}"
+    except Exception as e:
+        return f"Error: {str(e)}"
 
 with gr.Blocks() as demo:
-    gr.Markdown("# Cat vs Dog Classifier")
+    gr.Markdown("# 🐶🐱 Cat vs Dog Classifier")
+
     with gr.Row():
         with gr.Column():
-            image_input = gr.Image(shape=(224, 224))
-            classify_button = gr.Button("Classify")
+            image_input = gr.Image(
+                type="filepath", 
+                label="Input"
+            )
+            classify_button = gr.Button("🔍 Classify")
         with gr.Column():
-            output_text = gr.Textbox(label="Prediction")
-
-    classify_button.click(fn=classify_image, inputs=image_input, outputs=output_text)
+            output_text = gr.Textbox(label="🧠 Prediction", placeholder="Result will appear here")
 
-# Image input and output example
-demo = gr.Interface(
-    fn=greet, 
-    inputs=gr.inputs.Image(shape=(224, 224)), 
-    outputs="text"
-)
+    classify_button.click(fn=classify_image, inputs=[image_input], outputs=[output_text])
 
 demo.launch(debug=True)
-demo.launch(share=True)
+

src/config.py

@@ -35,6 +35,10 @@ class CatDogClassifierConfigs(BaseModel):
         default=2, 
         description="Number of output classes (cat and dog)"
     )
+    use_amp: bool = Field(
+        default=False,
+        description="Whether to use Automatic Mixed Precision (AMP) for training"
+    )
 
 class CatDogDatasetConfigsInput(BaseModel):
     data_path: str = Field(

src/data_preprocessing.py

@@ -22,6 +22,7 @@ class DataPreprocessor:
             transform = transforms.Compose([
                 transforms.Resize((self.image_size, self.image_size)),
                 transforms.RandomHorizontalFlip(p=self.horizontal_flip_prob),
+                transforms.RandomRotation(degrees=15),
                 transforms.ToTensor(),
                 transforms.Normalize(mean=self.mean, std=self.std),
             ])

src/infer.py

@@ -3,7 +3,8 @@ from src.model import CatDogClassifier
 from src.config import CatDogClassifierConfigs
 
 def inference_pipeline(
-       image_path: str = "datasets/single_prediction/cat_or_dog_1.jpg"
+    image_path: str = "datasets/single_prediction/cat_or_dog_1.jpg",
+    model_path: str = "checkpoints/ckpt_23_10_2025/best_cat_dog_classifier_model_20251019_122336.pth"
 ):
 
     # Initialize model
@@ -15,11 +16,12 @@ def inference_pipeline(
         kernel_size=3,
         stride=2,
         padding=1,
-        num_layers=3
+        num_layers=3,
+        use_amp=False
     )
     # Load state_dict
     model = CatDogClassifier(configs=model_configs)
-    model.load_state_dict(torch.load("cat_dog_classifier.pth", map_location="cpu"))
+    model.load_state_dict(torch.load(model_path, map_location="cpu"))
     y_pred = model.predict(
         model=model, 
         image_path=image_path
@@ -31,5 +33,5 @@ def inference_pipeline(
 
 
 if __name__ == "__main__":
-    y_pred = inference_pipeline("datasets/single_prediction/cat_or_dog_1.jpg")
+    y_pred = inference_pipeline("D:\\Desktop\\stores\\Application\\GoldenOwl\\technical_test\\test_image_2.jpg")
     print(y_pred)

src/model.py

@@ -1,8 +1,11 @@
-from tqdm.auto import tqdm
-from matplotlib import transforms
+
+from datetime import datetime
 import torch 
 import torchvision
 import torch.nn as nn
+
+from tqdm.auto import tqdm
+from torchvision import transforms
 from .config import CatDogClassifierConfigs
 
 class CatDogClassifier(nn.Module):
@@ -16,6 +19,8 @@ class CatDogClassifier(nn.Module):
         self.learning_rate = configs.learning_rate
         self.num_classes = configs.num_classes
         self.input_channels = configs.input_channels
+        self.device = configs.device
+        self.use_amp = configs.use_amp
 
         # Initialize the model architecture
         self._build_model()
@@ -37,28 +42,43 @@ class CatDogClassifier(nn.Module):
         self.conv_layer_2 = nn.Sequential(
             nn.Conv2d(
               in_channels=64, 
-              out_channels=512, 
+              out_channels=128, 
               kernel_size=self.kernel_size, 
               padding=self.padding
             ),
+            nn.BatchNorm2d(num_features=128),
             nn.ReLU(),
-            nn.BatchNorm2d(num_features=512),
             nn.MaxPool2d(kernel_size=2)
         )
         self.conv_layer_3 = nn.Sequential(
             nn.Conv2d(
-              in_channels=512, 
-              out_channels=512, 
+              in_channels=128, 
+              out_channels=256, 
               kernel_size=self.kernel_size, 
               padding=self.padding
             ),
+            nn.BatchNorm2d(num_features=256),
             nn.ReLU(),
-            nn.BatchNorm2d(num_features=512),
             nn.MaxPool2d(kernel_size=2)
         )
+        self.conv_layer_4 = nn.Sequential(
+            nn.Conv2d(
+              in_channels=256, 
+              out_channels=512, 
+              kernel_size=self.kernel_size, 
+              padding=self.padding
+            ),
+            nn.BatchNorm2d(num_features=512),
+            nn.ReLU(),
+            nn.AdaptiveAvgPool2d((1, 1))
+        )
         self.classifier = nn.Sequential(
           nn.Flatten(),
-          nn.Linear(in_features=512*3*3, out_features=self.num_classes)
+          nn.Dropout(p=0.5),
+          nn.Linear(in_features=512, out_features=256),
+          nn.ReLU(),
+          nn.Dropout(p=0.3),
+          nn.Linear(in_features=256, out_features=self.num_classes)
         )
 
 
@@ -66,9 +86,7 @@ class CatDogClassifier(nn.Module):
         x = self.conv_layer_1(x)
         x = self.conv_layer_2(x)
         x = self.conv_layer_3(x)
-        x = self.conv_layer_3(x)
-        x = self.conv_layer_3(x)
-        x = self.conv_layer_3(x)
+        x = self.conv_layer_4(x)
         x = self.classifier(x)
         return x
     
@@ -80,9 +98,14 @@ class CatDogClassifier(nn.Module):
             num_epochs: int,
             loss_fn: nn.Module,
             optimizer: torch.optim.Optimizer,
+            scheduler: torch.optim.lr_scheduler._LRScheduler = None,
     ):
-        # Placeholder for training logic
+        # Initialize the loss function and optimizer
+        scaler = torch.amp.GradScaler(device=self.configs.device, enabled=self.configs.use_amp)
+
         print("Training the model with provided data")
+        best_acc = 0.0
+
         # Implement training loop here
         results = {
             "train_loss": [],
@@ -98,13 +121,26 @@ class CatDogClassifier(nn.Module):
                 dataloader=train_dataloader,
                 loss_fn=loss_fn,
                 optimizer=optimizer,
+                epoch=epoch,
+                num_epochs=num_epochs,
+                scaler=scaler
             )
             test_loss, test_acc = self._test_step(
                 model=model,
                 dataloader=test_dataloader,
-                loss_fn=loss_fn,
+                loss_fn=loss_fn
             )
 
+            # ----- Scheduler update -----
+            if scheduler:
+                scheduler.step()
+
+            # ----- Save best model -----
+            if test_acc > best_acc:
+                best_acc = test_acc
+                torch.save(model.state_dict(), f"best_cat_dog_classifier_model_{datetime.now().strftime('%Y%m%d_%H%M%S')}.pth")
+
+
             results["train_loss"].append(train_loss)
             results["train_acc"].append(train_acc)
             results["test_loss"].append(test_loss)
@@ -115,7 +151,7 @@ class CatDogClassifier(nn.Module):
                 f"Train Loss: {train_loss:.4f}, Train Acc: {train_acc:.4f} | "
                 f"Test Loss: {test_loss:.4f}, Test Acc: {test_acc:.4f}"
             )
-
+        print(f"\n✅ Training complete! Best Test Accuracy: {best_acc:.4f}")
         return results
 
     def _train_step(
@@ -124,33 +160,43 @@ class CatDogClassifier(nn.Module):
             dataloader: torch.utils.data.DataLoader,
             loss_fn: nn.Module,
             optimizer: torch.optim.Optimizer,
+            epoch: int,
+            num_epochs: int,
+            scaler: torch.amp.GradScaler,
     ):
         # Define model in training mode
         model.train()
 
-        train_loss, train_acc = 0, 0
+        train_loss, train_acc, correct, total_train_examples = 0, 0, 0, 0
 
         # Loop through each batch
-        for batch_idx, (data, target) in enumerate(dataloader):
+        pbar = tqdm(enumerate(dataloader), desc=f"Epoch [{epoch+1}/{num_epochs}]")
+        for batch_idx, (data, target) in pbar:
             data, target = data.to(self.configs.device), target.to(self.configs.device)
 
             # print(f"Batch {batch_idx+1}: data shape {data.shape}, target shape {target.shape}")
             # Forward pass
-            y_pred = model(data)
-            # Calculate and accumulate loss
-            loss = loss_fn(y_pred, target)
+            # y_pred = model(data)
+            with torch.amp.autocast(device_type=self.configs.device, enabled=self.configs.use_amp):
+                y_pred = model(data)
+                # Calculate and accumulate loss
+                loss = loss_fn(y_pred, target)
+
             train_loss += loss.item()
             # Backward pass
             optimizer.zero_grad()
-            loss.backward()
-            optimizer.step()
-            # Calculate and accumulate accuracy metric across all batches
-            y_pred_labels = torch.argmax(torch.softmax(y_pred, dim=1), dim=1)
-            train_acc += (y_pred_labels == target).sum().item()/data.size(0)
+            scaler.scale(loss).backward()
+            scaler.step(optimizer)
+            scaler.update()
+
+            # Calculate and accumulate accuracy metric
+            y_pred_labels = torch.argmax(y_pred, dim=1)
+            correct += (y_pred_labels == target).sum().item()
+            total_train_examples += target.size(0)
 
         # Adjust loss and accuracy to get average loss and accuracy based on number of batches
         train_loss /= len(dataloader)
-        train_acc /= len(dataloader)
+        train_acc = correct / total_train_examples
 
         return train_loss, train_acc
 
@@ -162,8 +208,8 @@ class CatDogClassifier(nn.Module):
     ):
         # Define model in evaluation
         model.eval()
-        test_loss, test_acc = 0, 0
-        with torch.no_grad():
+        test_loss, test_acc, correct, total_test_examples = 0, 0, 0, 0
+        with torch.inference_mode():
             for batch_idx, (data, target) in enumerate(dataloader):
                 data, target = data.to(self.configs.device), target.to(self.configs.device)
                 # Forward pass
@@ -171,12 +217,13 @@ class CatDogClassifier(nn.Module):
                 # Calculate and accumulate loss
                 loss = loss_fn(y_pred, target)
                 test_loss += loss.item()
-                # Calculate and accumulate accuracy metric across all batches
-                y_pred_labels = torch.argmax(torch.softmax(y_pred, dim=1), dim=1)
-                test_acc += (y_pred_labels == target).sum().item()/data.size(0)
+                # Calculate and accumulate accuracy metric 
+                y_pred_labels = torch.argmax(y_pred, dim=1)
+                correct += (y_pred_labels == target).sum().item()
+                total_test_examples += target.size(0)
         # Adjust loss and accuracy to get average loss and accuracy based on number of batches
         test_loss /= len(dataloader)
-        test_acc /= len(dataloader)
+        test_acc = correct / total_test_examples
 
         return test_loss, test_acc
         
@@ -188,21 +235,21 @@ class CatDogClassifier(nn.Module):
     ) -> str:
         # Load and preprocess the image converting it to a tensor 
         # and normalizing the pixel values between 0 and 1
-        image_tensor = torchvision.io.read_image(str(image_path)).type(torch.float32) / 255.0
-        image_tensor_transformed = transforms.Compose([
-            transforms.Resize((256, 256)),
-            transforms.CenterCrop((224, 224)),
-            transforms.ToTensor(),
-        ])(image_tensor)
-
+        image_tensor = torchvision.io.read_image(str(image_path)).float() / 255.0
+        transform = transforms.Compose([
+            transforms.Resize((224, 224)),
+            transforms.Normalize(
+                mean=[0.485, 0.456, 0.406],
+                std=[0.229, 0.224, 0.225]
+            ),
+        ])
+        image_tensor_transformed = transform(image_tensor).unsqueeze(0).to(self.configs.device)
+        # Set model to evaluation mode and make prediction
+        model = model.to(self.configs.device)
         model.eval()
-        with torch.no_grad():
-            image_tensor_transformed = image_tensor_transformed.unsqueeze(0)  # Add batch dimension
-            image_tensor_pred = model(image_tensor_transformed).to(self.configs.device)
-            predicted_label = torch.argmax(torch.softmax(image_tensor_pred, dim=1), dim=1).item()
-
-            if predicted_label == 0:
-                return "cat"
-            else:
-                return "dog"
+        with torch.inference_mode():
+            image_tensor_pred = model(image_tensor_transformed)
+            predicted_label = torch.argmax(image_tensor_pred, dim=1).item()
+
+        return "cat" if predicted_label == 0 else "dog"
         

src/train.py

@@ -1,8 +1,8 @@
 import os
 import torch
 import torch.nn as nn
+from datetime import datetime
 
-from src import model
 from src.config import (
     CatDogDatasetConfigsInput,
     CatDogClassifierConfigs,
@@ -50,6 +50,7 @@ def train_pipeline():
     preprocessor = DataPreprocessor(data_preprocessing_configs)
     train_dataloader, test_dataloader = preprocessor.create_dataloader()
     device = "cuda" if torch.cuda.is_available() else "cpu"
+    print(f"Using device: {device}")
 
     # Model training
     model_configs = CatDogClassifierConfigs(
@@ -60,7 +61,8 @@ def train_pipeline():
         kernel_size=3,
         stride=2,
         padding=1,
-        num_layers=3
+        num_layers=3,
+        use_amp=True
     )
 
     model = CatDogClassifier(model_configs)
@@ -81,7 +83,7 @@ def train_pipeline():
         model=model,
         train_dataloader=train_dataloader,
         test_dataloader=test_dataloader,
-        num_epochs=30,
+        num_epochs=20,
         loss_fn=loss_fn,
         optimizer=optimizer
     )
@@ -89,7 +91,7 @@ def train_pipeline():
     print(f"Training completed in {end_time - start_time} seconds.")
 
     # Save the trained model
-    torch.save(model.state_dict(), "cat_dog_classifier.pth")
+    torch.save(model.state_dict(), f"cat_dog_classifier_{datetime.now().strftime('%Y%m%d_%H%M%S')}.pth")
     print("Model saved to cat_dog_classifier.pth")