♻️ CNN Garbage Classification using PyTorch

This repository contains a Deep Learning pipeline for classifying images of waste into 6 distinct categories. The project leverages Transfer Learning using a pre-trained ResNet50 model built with PyTorch.

📋 Project Overview

The goal of this project is to automatically classify images of garbage to assist in recycling and waste management.

The model is trained on a dataset of 2,000 labeled images and is designed to predict the classes of 500 unlabeled test images, generating a formatted submission.csv file.

🏷️ Classes

The dataset is divided into 6 categories, mapped to numeric IDs:

• 0: paper
• 1: glass
• 2: plastic
• 3: metal
• 4: cardboard
• 5: trash

🧠 Model Architecture

• Base Model: ResNet50 (Pre-trained on ImageNet1K)
• Transfer Learning Strategy: The base convolutional layers are frozen to retain feature extraction capabilities. The final Fully Connected (fc) layer is replaced with a new linear layer (nn.Linear(2048, 6)) to output predictions for our specific 6 classes.
• Loss Function: Cross-Entropy Loss
• Optimizer: Adam (Learning Rate = 0.001)

📁 Repository Structure

Ensure your local directory matches this structure before running the code:

CnnGarbageClassification/
│
├── Dataset/
│   ├── Train/              # Folder containing the 2000 labeled images
│   ├── Test/               # Folder containing the 500 unlabeled test images
│   └── Train_labels.csv    # CSV mapping Train images to their text labels
│
├── WasteDataset.py         # Custom PyTorch Dataset class
├── main.py                 # Main script (Data processing, Training, Inference)
└── README.md               # Project documentation

⚙️ Installation & Requirements

This project requires Python 3.8+. It is highly recommended to run this on a machine with a CUDA-enabled GPU (e.g., NVIDIA RTX 3060 Ti) for significantly faster training.

git clone [https://github.com/YacIneTheBox/CnnGarbageClassification.git](https://github.com/YacIneTheBox/CnnGarbageClassification.git)

cd CnnGarbageClassification

pip install pandas scikit-learn pillow matplotlib

pip install torch torchvision torchaudio --index-url [https://download.pytorch.org/whl/cu121](https://download.pytorch.org/whl/cu121)

python main.py

👾 What the script does:

• Data Prep: Loads Train_labels.csv, cleans the data, maps text labels to numeric IDs, and performs a stratified 80/20 train-validation split.

• Training: Trains the ResNet50 model over 5 epochs (customizable) and evaluates it on the validation set.

• Inference: Loads the 500 unlabelled images from Dataset/Test.

• Submission: Generates a submission.csv file formatted strictly as Id,label (Image filename, Predicted numeric ID) ready for competition upload.

📊 Performance Baseline:

With basic configurations and 5 to 10 epochs, the model achieves an approximate accuracy of:

• Training Accuracy: ~89 - 91%

• Validation Accuracy: ~83 - 85%

🛠️ Future Improvements:

To push the validation accuracy beyond 90%, the following techniques can be implemented:

• Data Augmentation: Apply RandomCrop, RandomHorizontalFlip, and RandomRotation to the training transforms to combat overfitting.

• Class Weights: Implement weighted Cross-Entropy Loss to handle the heavy imbalance of the trash class.

• Fine-Tuning: Unfreeze the deeper convolutional layers of the ResNet50 model and train with a very low learning rate.

🎲 The Dataset Link:

https://www.kaggle.com/competitions/garbage-classification-challenge/overview