Pig Counting Model

Pig counter model for the SCAI Is No Limit competition.

This project implements a simple deep learning model in PyTorch to count the number of pigs in an image using a supervised regression approach. The model is lightweight, end-to-end, and includes visualization of predictions via saliency heatmaps.

Dataset

This project uses the Pig Counting Challenge Dataset from Kaggle:

https://www.kaggle.com/datasets/startears/pig-counting-challenge

Features

• Custom dataset loader for pig images and label files
• Convolutional Neural Network (CNN) regression model for pig count
• Data augmentation using torchvision transforms
• Training, testing, and evaluation loops
• Mean Squared Error (MSE) loss evaluation
• Saliency visualization to highlight model focus areas
• Dataset reshuffling script for creating train/test splits

How It Works

This project trains a CNN to predict the number of pigs in an image: Images and label text files are paired (labels count the number of pigs). Each line in the .txt file represents one pig. The training pipeline loads images and predicts a single scalar output (pig count). The model is trained using MSE loss. After training, the model generates saliency heatmaps for interpretability. The code is implemented in Python using PyTorch.

Getting Started

Clone the Repo

git clone https://github.com/reddykingit/Pig-Counting-Model.git

cd Pig-Counting-Model

#Project Structure

Pig-Counting-Model/
├── images/             # Images split into train/test
├── labels/             # Corresponding label files (.txt)
├── main.py             # Train / test pipeline + visualization
├── pigdatasetloader.py # Custom dataset loader
├── reshuffle.py        # Script to reshuffle dataset
├── visualization.py    # Saliency visualization functions
├── custom_image.jpg    # Optional image for custom testing
└── .gitignore

Dependencies

pip install torch torchvision matplotlib pillow

Training the Model

Before training:

Place your images under:

images/train
images/test

Place corresponding .txt label files under:

labels/train
labels/test

Then run:

python main.py

This will:

• Train the model
• Output training and test performance
• Print MSE loss and average prediction error
• Generate saliency visualizations

Dataset Shuffle

If you want to randomly split a dataset into training and test sets:

python reshuffle.py

WARNING: This moves files in your project directories — use on a copy of your dataset.

Visualization

The model creates heatmap overlays showing which pixels most influenced the count prediction. These visualizations are displayed using matplotlib when running main.py.

Example Output

During training/testing, the console prints:

Using device: cuda (or cpu)
Epoch X/X, Average Loss: ...
On average, the model is off by: Y.ZZ pigs
Test Set MSE Loss: ...
Generating Heatmaps...

An example of the visual comparison of prediction vs actual count is shown in a plotted window.

Notes

• Input images are resized to 256×256.
• The model architecture is a simple CNN regression head.
• The architecture can be extended with object detection models (e.g., YOLO) or density map approaches.

Acknowledgements

This project won 1st place in the SCAI Is No Limit Competition and uses the Pig Counting Challenge dataset from Kaggle.