End-to-End Computer Vision Pipeline: Binary Classification

An end-to-end Machine Learning Engineering pipeline for binary image classification (Cats vs. Dogs). This project demonstrates the complete ML lifecycle, moving beyond simple notebooks to establish a robust, production-ready codebase.

🚀 Architecture Overview

This repository implements a modular pipeline with the following capabilities:

• Data Engineering: Automated ingestion, shuffling, integrity validation (corruption catching), and train/val/test splitting.
• Data Streaming: Highly optimized PyTorch DataLoaders with asynchronous fetching and stochastic data augmentation.
• Model Training: Transfer learning leveraging pre-trained CNN architectures (e.g., ResNet).
• Experiment Tracking: Integration with MLflow for metric logging and model registry.
• Serving & Portability: Model export to standard ONNX format, served via a FastAPI REST endpoint with a custom modern Web UI.

📁 Repository Structure

cv-classifier/
├── data/
│   ├── raw/             # Raw dataset (Ignored by Git)
│   └── processed/       # Validated/split data ready for PyTorch (Ignored by Git)
├── img/                 # Documentation images
├── models/              # Serialized model weights and ONNX exports (Ignored by Git)
├── src/
│   ├── data_prep.py     # Data validation, shuffling, and partitioning script
│   ├── train.py         # Transforms, DataLoaders, and Training Loop
│   ├── export.py        # PyTorch to ONNX conversion logic
│   └── serve.py         # FastAPI REST endpoint and UI for inference
├── .gitignore
├── requirements.txt
└── README.md

🛠️ Environment Setup

This project requires Python 3.11 and PyTorch. It is heavily recommended to run this pipeline with NVIDIA CUDA-enabled hardware for accelerated training.

1. Clone the repository:

   git clone [https://github.com/gerymeee/cv-classifier.git](https://github.com/gerymeee/cv-classifier.git)
   cd cv-classifier

2. Create and activate a virtual environment (Ensure you are using Python 3.11):

   # Windows (using the Python launcher)
   py -3.11 -m venv .venv
   .venv\Scripts\activate
   
   # Mac/Linux
   python3.11 -m venv .venv
   source .venv/bin/activate

3. Install dependencies: (Note: The requirements file assumes a CUDA 11.8 environment. Modify the PyTorch index URL in requirements.txt if using CPU or a different CUDA version).

   pip install -r requirements.txt

📊 Dataset Ingestion

Note: Because this pipeline is a robust binary classifier, it is not strictly limited to Cats vs. Dogs. You can supply entirely different image datasets (e.g., Pizza vs. Hotdog, Cars vs. Bikes) tailored to your specific needs by simply changing the source data.

1. Place your unorganized binary classification images into the raw directory, separated by class folders (e.g., data/raw/Cat/ and data/raw/Dog/).

2. Run the data preparation script to clean and partition the data.

   python src/data_prep.py

🚂 Running the Pipeline

Run the full training pipeline (includes interactive DataLoader sanity check, Transfer Learning via ResNet-18, and MLflow tracking):

```bash
python src/train.py
```

Export the trained PyTorch weights to a highly optimized ONNX graph for serving:

```bash
python src/export.py
```

🌐 Model Serving & Web UI

Launch the FastAPI server to interact with your trained model via a modern, responsive web interface:

```bash
python src/serve.py
```

Navigate to http://127.0.0.1:8000/ in your browser to access the radar.