Image Classification Using CNN

This project contains a simple convolutional neural network for image classification with PyTorch. The included example trains a binary classifier for cats and dogs and can also run inference on a single image after training.

Features

• Builds a grayscale image dataset for training
• Downloads and extracts the Cats vs Dogs dataset automatically
• Splits the data into training and validation sets
• Trains a convolutional neural network with fully connected layers
• Saves model weights and optimizer state for reuse
• Runs prediction on a single image such as cat.jpg

Project Structure

• catvsdog_example.py: main example script for training and prediction
• image_classifier/classifier.py: training loop, validation, saving, loading, and prediction logic
• image_classifier/cnn_model.py: neural network architecture definition
• image_classifier/tools.py: dataset preparation, image loading, plotting, and helper functions
• cat.jpg: sample image for prediction

Requirements

Install the required Python packages before running the example:

• torch
• torchvision
• tqdm
• opencv-python
• wget
• matplotlib
• numpy

Example install command:

pip install torch torchvision tqdm opencv-python wget matplotlib numpy

How It Works

When catvsdog_example.py runs with parameters['rebuild'] = True, the script:

1. Downloads the dataset archive
2. Extracts the image folders
3. Resizes images to the configured resolution
4. Converts images to grayscale
5. Normalizes pixel values
6. Saves the processed dataset as dataset.npy
7. Trains the model and stores outputs in the configured save directory

If parameters['rebuild'] = False, the script loads a previously saved model instead of rebuilding the dataset and retraining.

Running the Project

Run the example with:

python catvsdog_example.py

The current configuration in the example script:

• uses images from PetImages/Cat and PetImages/Dog
• resizes images to 50 x 50
• trains for 6 epochs
• uses a validation split of 10%
• saves outputs to CNN_catsVsdogs
• predicts the class of cat.jpg

Output Files

After training, the script stores generated artifacts in the save directory defined by parameters['savedir']. With the current configuration, that directory is CNN_catsVsdogs and typically contains:

• the saved model file
• the processed dataset file
• the extracted PetImages folder

Configuration

The main configuration is stored in the parameters dictionary inside catvsdog_example.py. Important values you may want to change include:

• parameters['size']: image resolution
• parameters['device']: execution device, currently set to cuda
• parameters['epochs']: number of training epochs
• parameters['learning rate']: optimizer learning rate
• parameters['batch_size_training']: training batch size
• parameters['batch_size_validation']: validation batch size
• parameters['savedir']: directory used for model outputs
• parameters['dataset']: processed dataset filename

Using CPU Instead of GPU

The current example is configured with:

parameters['device'] = 'cuda'

If your machine does not have CUDA available, change it to:

parameters['device'] = 'cpu'

Adapting the Project

You can adapt this repository to another image-classification problem by changing the dataset paths and labels, then adjusting the output size and prediction logic as needed. For binary image classification, the current structure is already close to reusable.

Notes

• The example currently uses grayscale images
• The project is designed for binary classification in its current form
• Training downloads data from an external source, so internet access is required when rebuilding the dataset

Author

Aditi Sah