upload all files

.gitignore

@@ -205,3 +205,10 @@ cython_debug/
 marimo/_static/
 marimo/_lsp/
 __marimo__/
+
+
+# local
+datasets/*
+checkpoints/*
+dc_env/*
+note.md 

README.md

@@ -0,0 +1,27 @@
+##  Image classification - Cat & Dog Classification
+
+###    Prerequistiion Requirements:
+- Training Processing:
+    - Datasets Zip file to be saved in datasets folder.
+- Dependencies Installation:
+    - Create an virtual environment with conda
+        - conda `create -p dc_env python=3.9 -y`
+        - Activate created env: `conda activate dc_env/`
+    - Using `pip install -r requirements.txt`
+- Other Requirements:
+    - GPU (E.g: NVIDIA RTX 3050,...)
+    - Python version >= 3.9
+
+## How to run this project:
+### Training the model:
+- Utilizing `python src/train.py`
+
+### Run Deployment on your local:
+- Utilizing `python deployment/gradio/main.py` \
+Or
+- On HuggingFace Server, you can access at: `` 
+
+
+
+
+

deployment/gradio/main.py

@@ -0,0 +1,30 @@
+import gradio as gr
+
+
+def greet(name):
+    return f"Hello {name}!"
+
+def classify_image(image):
+    return "cat"  # Placeholder for actual image classification logic
+
+
+with gr.Blocks() as demo:
+    gr.Markdown("# Cat vs Dog Classifier")
+    with gr.Row():
+        with gr.Column():
+            image_input = gr.Image(shape=(224, 224))
+            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)
+
+# Image input and output example
+demo = gr.Interface(
+    fn=greet, 
+    inputs=gr.inputs.Image(shape=(224, 224)), 
+    outputs="text"
+)
+
+demo.launch(debug=True)
+demo.launch(share=True)

requirements.txt

@@ -0,0 +1,4 @@
+scikit-learn
+gradio
+pydantic
+torch torchvision --index-url https://download.pytorch.org/whl/cu126

src/config.py

@@ -0,0 +1,94 @@
+from pydantic import BaseModel, Field
+from typing import List, Dict, Optional
+
+
+class CatDogClassifierConfigs(BaseModel):
+    device: str = Field(
+        default="cpu",
+        description="Device to run the model on (cpu or cuda)"
+    )
+    input_channels: int = Field(
+        default=3, 
+        description="Number of input channels for the images"
+    )
+    kernel_size: int = Field(
+        default=3, 
+        description="Size of the convolutional kernel"
+    )
+    stride: int = Field(
+        default=1, 
+        description="Stride for the convolutional layers"
+    )
+    padding: int = Field(
+        default=1, 
+        description="Padding for the convolutional layers"
+    )
+    num_layers: int = Field(
+        default=2, 
+        description="Number of convolutional layers"
+    )
+    learning_rate: float = Field(
+        default=0.001, 
+        description="Learning rate for the optimizer"
+    )
+    num_classes: int = Field(
+        default=2, 
+        description="Number of output classes (cat and dog)"
+    )
+
+class CatDogDatasetConfigsInput(BaseModel):
+    data_path: str = Field(
+        default="datasets/datasets.zip",
+        description="Path to the dataset (can be a folder or an archive file)"
+    )
+    train_data_path: Optional[str] = Field(
+        default="datasets/train", 
+        description="Path to the training data"
+    )
+    test_data_path: Optional[str] = Field(
+        default="datasets/test", 
+        description="Path to the testing data"
+    )
+    test_size: Optional[float] = Field(
+        default=0.2,
+        description="Proportion of the dataset to include in the test split"
+    )
+    random_state: Optional[int] = Field(
+        default=42,
+        description="Random seed for data splitting"
+    )
+
+class DataPreprocessorConfigsInput(BaseModel):
+    train_dataset_path: str = Field(
+        default="datasets/train", 
+        description="Path to the training dataset"
+    )
+    test_dataset_path: str = Field(
+        default="datasets/test", 
+        description="Path to the testing dataset"
+    )
+    shuffle: bool = Field(
+        default=True,
+        description="Whether to shuffle the data during loading"
+    )
+    batch_size: int = Field(
+        default=64, 
+        description="Number of samples per batch"
+    )
+    horizontal_flip_prob: float = Field(
+        default=0.5,
+        description="Probability of applying random horizontal flip"
+    )
+
+    image_size: int = Field(
+        default=224, 
+        description="Size to which images will be resized"
+    )
+    mean: List[float] = Field(
+        default=[0.485, 0.456, 0.406], 
+        description="Mean for normalization"
+    )
+    std: List[float] = Field(
+        default=[0.229, 0.224, 0.225], 
+        description="Standard deviation for normalization"
+    )

src/data_ingestion.py

@@ -0,0 +1,87 @@
+import os
+import shutil
+import zipfile
+
+from sklearn.model_selection import train_test_split
+from .config import CatDogDatasetConfigsInput
+
+class CatDogDataset:
+    def __init__(self, data_configs: CatDogDatasetConfigsInput):
+        self.configs = data_configs
+        self.data_path = data_configs.data_path
+        self.train_data_path = data_configs.train_data_path
+        self.test_data_path = data_configs.test_data_path
+        self.test_size = data_configs.test_size if hasattr(data_configs, 'test_size') else 0.2
+        self.random_state = data_configs.random_state if hasattr(data_configs, 'random_state') else 42
+
+        self.file_type = self.data_path.split('.')[-1]
+
+    # Check the type of the data path input
+    def _check_type(self):
+        if self.file_type in ['zip', 'tar', 'tar.gz']:
+            return "archive"
+        elif self.file_type == '':
+            return "folder"
+        else:
+            raise ValueError(f"Unsupported file type: {self.file_type}")
+        
+    # Extract archive files if data path is an archive
+    def _extract_archive(self):
+        print(f"Extracting archive: {self.data_path}")
+        extract_dir = os.path.splitext(self.data_path)[0]
+        os.makedirs(extract_dir, exist_ok=True)
+        with zipfile.ZipFile(self.data_path, 'r') as zip_ref:
+            zip_ref.extractall(extract_dir)
+            print(f"Extracted archive to {extract_dir}")
+        
+        # Remove the original archive file after extraction
+        os.remove(self.data_path)
+        print(f"Removed archive file: {self.data_path}")
+        # Update data_path to point to the extracted folder
+        self.data_path = self.data_path.rstrip('.zip').rstrip('.tar').rstrip('.gz')
+
+
+
+    def _split_data(self):
+        # Split the original dataset into training and testing sets folder
+        try:
+            # Run into each folder (cats and dogs) and split the images
+            for c in os.listdir(self.data_path):
+                all_images = os.listdir(os.path.join(self.data_path, c))
+                train_images, test_images = train_test_split(
+                    all_images, 
+                    test_size=self.test_size, 
+                    random_state=self.random_state
+                )
+
+                # Create train and test directories if they don't exist
+                os.makedirs(os.path.join(self.train_data_path, c), exist_ok=True)
+                os.makedirs(os.path.join(self.test_data_path, c), exist_ok=True)
+
+                # Move images to respective folders
+                for img in train_images:
+                    shutil.move(os.path.join(self.data_path, c, img), os.path.join(self.train_data_path, c, img))
+                for img in test_images:
+                    shutil.move(os.path.join(self.data_path, c, img), os.path.join(self.test_data_path, c, img))
+
+            print("Data split successfully")
+            # Remove the original data folder after splitting
+            shutil.rmtree(self.data_path)
+            print("Original data folder removed")
+
+
+        except Exception as e:
+            print(f"Error splitting data: {e}")
+
+        
+    def load_data(self):
+        # Logic to load and preprocess the dataset
+        data_type = self._check_type()
+        if data_type == "archive":
+            self._extract_archive()
+        elif data_type == "folder":
+            print(f"Loading data from folder: {self.data_path}")
+        else:
+            raise ValueError("Unsupported data type")
+        self._split_data()
+        print("Data loading and preprocessing completed")

src/data_preprocessing.py

@@ -0,0 +1,57 @@
+import torch 
+from torch.utils.data import Dataset, DataLoader
+from torchvision import datasets, transforms
+
+from .config import DataPreprocessorConfigsInput
+
+class DataPreprocessor:
+    def __init__(self, configs: DataPreprocessorConfigsInput):
+        self.configs = configs
+        self.train_dataset_path = self.configs.train_dataset_path
+        self.test_dataset_path = self.configs.test_dataset_path
+        self.shuffle = self.configs.shuffle
+        self.batch_size = self.configs.batch_size
+        self.horizontal_flip_prob = self.configs.horizontal_flip_prob  # Probability for random horizontal flip
+        self.image_size = self.configs.image_size
+        self.mean = self.configs.mean
+        self.std = self.configs.std
+
+
+    def preprocess(self, label:str = "train") -> transforms.Compose:
+        if label == "train":
+            transform = transforms.Compose([
+                transforms.Resize((self.image_size, self.image_size)),
+                transforms.RandomHorizontalFlip(p=self.horizontal_flip_prob),
+                transforms.ToTensor(),
+                transforms.Normalize(mean=self.mean, std=self.std),
+            ])
+
+        else:
+            transform = transforms.Compose([
+                transforms.Resize((self.image_size, self.image_size)),
+                transforms.ToTensor(),
+                transforms.Normalize(mean=self.mean, std=self.std)
+            ])
+        return transform 
+
+    def create_dataloader(self) -> DataLoader:
+        train_transform = self.preprocess()
+        test_transform = self.preprocess(label="test")
+        train_dataset = datasets.ImageFolder(root=self.train_dataset_path, transform=train_transform)
+        train_dataloader = DataLoader(
+            dataset=train_dataset, 
+            batch_size=self.batch_size,
+            shuffle=self.shuffle
+        )
+
+        test_dataset = datasets.ImageFolder(root=self.test_dataset_path, transform=test_transform)
+        test_dataloader = DataLoader(
+            dataset=test_dataset, 
+            batch_size=self.batch_size,
+            shuffle= not self.shuffle
+        )
+
+        print(f"Train dataset size: {len(train_dataset)}")
+        print(f"Test dataset size: {len(test_dataset)}")
+
+        return train_dataloader, test_dataloader

src/infer.py

@@ -0,0 +1,35 @@
+import torch
+from src.model import CatDogClassifier
+from src.config import CatDogClassifierConfigs
+
+def inference_pipeline(
+       image_path: str = "datasets/single_prediction/cat_or_dog_1.jpg"
+):
+
+    # Initialize model
+    model_configs = CatDogClassifierConfigs(
+        device="cuda" if torch.cuda.is_available() else "cpu",
+        input_channels=3,
+        num_classes=2,
+        learning_rate=0.001,
+        kernel_size=3,
+        stride=2,
+        padding=1,
+        num_layers=3
+    )
+    # Load state_dict
+    model = CatDogClassifier(configs=model_configs)
+    model.load_state_dict(torch.load("cat_dog_classifier.pth", map_location="cpu"))
+    y_pred = model.predict(
+        model=model, 
+        image_path=image_path
+    )
+    print(f"Predicted class for the image {image_path}: {y_pred}")
+
+    return y_pred
+
+
+
+if __name__ == "__main__":
+    y_pred = inference_pipeline("datasets/single_prediction/cat_or_dog_1.jpg")
+    print(y_pred)

src/model.py

@@ -0,0 +1,208 @@
+from tqdm.auto import tqdm
+from matplotlib import transforms
+import torch 
+import torchvision
+import torch.nn as nn
+from .config import CatDogClassifierConfigs
+
+class CatDogClassifier(nn.Module):
+    def __init__(self, configs: CatDogClassifierConfigs):
+        super(CatDogClassifier, self).__init__()
+        self.configs = configs
+        self.kernel_size = configs.kernel_size
+        self.stride = configs.stride
+        self.padding = configs.padding
+        self.num_layers = configs.num_layers
+        self.learning_rate = configs.learning_rate
+        self.num_classes = configs.num_classes
+        self.input_channels = configs.input_channels
+
+        # Initialize the model architecture
+        self._build_model()
+
+
+    def _build_model(self):
+        # Placeholder for model building logic
+        self.conv_layer_1 = nn.Sequential(
+            nn.Conv2d(
+              in_channels=self.input_channels, 
+              out_channels=64, 
+              kernel_size=self.kernel_size, 
+              padding=self.padding
+            ),
+            nn.ReLU(),
+            nn.BatchNorm2d(num_features=64),
+            nn.MaxPool2d(kernel_size=2)
+        )
+        self.conv_layer_2 = nn.Sequential(
+            nn.Conv2d(
+              in_channels=64, 
+              out_channels=512, 
+              kernel_size=self.kernel_size, 
+              padding=self.padding
+            ),
+            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, 
+              kernel_size=self.kernel_size, 
+              padding=self.padding
+            ),
+            nn.ReLU(),
+            nn.BatchNorm2d(num_features=512),
+            nn.MaxPool2d(kernel_size=2)
+        )
+        self.classifier = nn.Sequential(
+          nn.Flatten(),
+          nn.Linear(in_features=512*3*3, out_features=self.num_classes)
+        )
+
+
+    def forward(self, x: torch.Tensor):
+        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.classifier(x)
+        return x
+    
+    def train_process(
+            self, 
+            model: nn.Module,
+            train_dataloader: torch.utils.data.DataLoader, 
+            test_dataloader: torch.utils.data.DataLoader,
+            num_epochs: int,
+            loss_fn: nn.Module,
+            optimizer: torch.optim.Optimizer,
+    ):
+        # Placeholder for training logic
+        print("Training the model with provided data")
+        # Implement training loop here
+        results = {
+            "train_loss": [],
+            "train_acc": [],
+            "test_loss": [],
+            "test_acc": []
+        }
+
+        # Loop through each epoch
+        for epoch in tqdm(range(num_epochs)):
+            train_loss, train_acc = self._train_step(
+                model=model,
+                dataloader=train_dataloader,
+                loss_fn=loss_fn,
+                optimizer=optimizer,
+            )
+            test_loss, test_acc = self._test_step(
+                model=model,
+                dataloader=test_dataloader,
+                loss_fn=loss_fn,
+            )
+
+            results["train_loss"].append(train_loss)
+            results["train_acc"].append(train_acc)
+            results["test_loss"].append(test_loss)
+            results["test_acc"].append(test_acc)
+
+            print(
+                f"Epoch [{epoch+1}/{num_epochs}] "
+                f"Train Loss: {train_loss:.4f}, Train Acc: {train_acc:.4f} | "
+                f"Test Loss: {test_loss:.4f}, Test Acc: {test_acc:.4f}"
+            )
+
+        return results
+
+    def _train_step(
+            self,
+            model: nn.Module,
+            dataloader: torch.utils.data.DataLoader,
+            loss_fn: nn.Module,
+            optimizer: torch.optim.Optimizer,
+    ):
+        # Define model in training mode
+        model.train()
+
+        train_loss, train_acc = 0, 0
+
+        # Loop through each batch
+        for batch_idx, (data, target) in enumerate(dataloader):
+            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)
+            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)
+
+        # Adjust loss and accuracy to get average loss and accuracy based on number of batches
+        train_loss /= len(dataloader)
+        train_acc /= len(dataloader)
+
+        return train_loss, train_acc
+
+    def _test_step(
+            self,
+            model: nn.Module,
+            dataloader: torch.utils.data.DataLoader,
+            loss_fn: nn.Module,
+    ):
+        # Define model in evaluation
+        model.eval()
+        test_loss, test_acc = 0, 0
+        with torch.no_grad():
+            for batch_idx, (data, target) in enumerate(dataloader):
+                data, target = data.to(self.configs.device), target.to(self.configs.device)
+                # Forward pass
+                y_pred = model(data)
+                # 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)
+        # Adjust loss and accuracy to get average loss and accuracy based on number of batches
+        test_loss /= len(dataloader)
+        test_acc /= len(dataloader)
+
+        return test_loss, test_acc
+        
+
+    def predict(
+            self, 
+            model: nn.Module,
+            image_path: str
+    ) -> 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)
+
+        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"
+        

src/train.py

@@ -0,0 +1,97 @@
+import os
+import torch
+import torch.nn as nn
+
+from src import model
+from src.config import (
+    CatDogDatasetConfigsInput,
+    CatDogClassifierConfigs,
+    DataPreprocessorConfigsInput
+)
+from src.data_ingestion import CatDogDataset
+from src.data_preprocessing import DataPreprocessor
+from src.model import CatDogClassifier
+
+
+def train_pipeline():
+    # Setup Intialize configurations
+    # Data ingestion
+    data_path = "datasets/datasets.zip"
+    train_data_path = "datasets/train"
+    test_data_path = "datasets/test"
+
+    data_path = os.path.join(os.getcwd(), data_path)
+    train_data_path = os.path.join(os.getcwd(), train_data_path)
+    test_data_path = os.path.join(os.getcwd(), test_data_path)
+
+
+    # data_ingestion_configs = CatDogDatasetConfigsInput(
+    #     data_path=data_path,
+    #     train_data_path=train_data_path,
+    #     test_data_path=test_data_path,
+    #     test_size=0.2,
+    #     random_state=42
+    # )
+    # print(data_ingestion_configs)
+    # dataset = CatDogDataset(data_ingestion_configs)
+    # dataset.load_data()
+    
+    # Data preprocessing
+    data_preprocessing_configs = DataPreprocessorConfigsInput(
+        train_dataset_path=train_data_path,
+        test_dataset_path=test_data_path,
+        shuffle=True,
+        batch_size=32,
+        horizontal_flip_prob=0.5,
+        image_size=224,
+        mean=[0.485, 0.456, 0.406],
+        std=[0.229, 0.224, 0.225]
+    )
+    preprocessor = DataPreprocessor(data_preprocessing_configs)
+    train_dataloader, test_dataloader = preprocessor.create_dataloader()
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+
+    # Model training
+    model_configs = CatDogClassifierConfigs(
+        device=device,
+        input_channels=3,
+        num_classes=2,
+        learning_rate=0.001,
+        kernel_size=3,
+        stride=2,
+        padding=1,
+        num_layers=3
+    )
+
+    model = CatDogClassifier(model_configs)
+    model.to(device)
+    # Set random seeds
+    torch.manual_seed(42) 
+    torch.cuda.manual_seed(42)
+    # Setup loss function and optimizer
+    loss_fn = nn.CrossEntropyLoss()
+    optimizer = torch.optim.Adam(params=model.parameters(), lr=0.001)
+
+    # # Calculate training time using timeit
+    # Start the timer
+    from timeit import default_timer as timer 
+    start_time = timer()
+
+    model.train_process(
+        model=model,
+        train_dataloader=train_dataloader,
+        test_dataloader=test_dataloader,
+        num_epochs=30,
+        loss_fn=loss_fn,
+        optimizer=optimizer
+    )
+    end_time = timer()
+    print(f"Training completed in {end_time - start_time} seconds.")
+
+    # Save the trained model
+    torch.save(model.state_dict(), "cat_dog_classifier.pth")
+    print("Model saved to cat_dog_classifier.pth")
+
+
+if __name__ == "__main__":
+    train_pipeline()

tests/data_ingestion.py

@@ -0,0 +1,17 @@
+from src.config import CatDogDatasetConfigsInput
+from src.data_ingestion import CatDogDataset
+def main():
+    # Initialize data ingestion with configurations
+    data_configs = CatDogDatasetConfigsInput(
+        data_path="datasets/datasets.zip",
+        train_data_path="datasets/train",
+        test_data_path="datasets/test",
+        test_size=0.2,
+        random_state=42
+    )
+    print(data_configs)
+    dataset = CatDogDataset(data_configs)
+    dataset.load_data()
+
+if __name__ == "__main__":
+    main()

tests/test_gpu.py

@@ -0,0 +1,6 @@
+import torch
+
+if __name__ == "__main__":
+    print(torch.cuda.is_available())
+    print(torch.version.cuda)
+    print(torch.cuda.get_device_name(0) if torch.cuda.is_available() else "No GPU detected")