pytorch-cnn-image-classifier

PyTorch CNN Image Classifier

A modular, production-ready Convolutional Neural Network (CNN) architecture implemented in PyTorch for multi-class image classification.

Features

• Modular PyTorch Architecture (model.py): 3 convolutional blocks with Batch Normalization, ReLU, MaxPool, Dropout, Adaptive Average Pooling, and a classification head.
• Ready-to-Use Inference Pipeline (predict.py): Predict directly on images or run instant synthetic demo tests with top-k ranked probabilities and CIFAR-10 labels.
• Automated Test Suite (test_pipeline.py): Unit tests validating forward pass, variable resolution handling, gradient backpropagation flow, and checkpoint saving/loading.
• Interactive Local Web App (app.py): Built-in Gradio interface to test and classify uploaded images in your browser.
• Config & Label Mapping (config.json): Formal model metadata with id2label mapping for the 10 classes.
• Training Script (train.py): Full training and validation pipeline on CIFAR-10 with data augmentation, Cosine Annealing learning rate scheduling, and automatic best-model checkpointing.

---

Quickstart

1. Installation

git clone https://huggingface.co/sarvesh74/pytorch-cnn-image-classifier
cd pytorch-cnn-image-classifier
pip install -r requirements.txt

2. Verify Installation & Architecture

Run the test suite to verify that all layers, tensor dimensions, and backward passes function as expected:

python test_pipeline.py

3. Run Inference

Instant Synthetic Demo

Test inference immediately without needing an external image file:

python predict.py --demo

Classify an Image File

python predict.py --image path/to/your/image.jpg --top-k 3

Python API Usage

from PIL import Image
from predict import ImageClassifier

classifier = ImageClassifier()
image = Image.open("your_image.jpg")
predictions = classifier.predict(image, top_k=3)

for pred in predictions:
    print(f"{pred['label']}: {pred['confidence']:.2f}%")

4. Launch the Interactive Web App (Gradio)

To launch an interactive browser UI where you can upload and test images:

python app.py

Open the provided local URL (typically http://127.0.0.1:7860) in your browser.

5. Training on CIFAR-10

Train the network from scratch using the included pipeline:

python train.py --epochs 20 --batch-size 64 --lr 0.001

The script will automatically download the dataset, track validation accuracy across epochs, and export the best checkpoint to best_model.pth.

---

Model Architecture Details

Spatial Dimension Formula

$$\text{Output Size} = \left\lfloor \frac{\text{Input} - \text{Kernel} + 2 \times \text{Padding}}{\text{Stride}} \right\rfloor + 1$$