Skip to main content

Finetunned ResNet18 for garbage image classification

Project description

๐Ÿ—‘๏ธ Garbage Classifier

Python PyTorch Lightning License

A deep learning-based waste classification system using ResNet18 and PyTorch Lightning

Features โ€ข Installation โ€ข Usage โ€ข Documentation โ€ข Results


๐Ÿ“‹ Overview

This project implements an image classification system to automatically categorize waste materials into six classes: cardboard, glass, metal, paper, plastic, and trash. Built with PyTorch Lightning and ResNet18, it provides a modular, scalable solution for waste management automation.

๐ŸŽฏ Key Features

  • โœจ ResNet18 backbone fine-tuned for garbage classification
  • โšก PyTorch Lightning framework for clean, scalable training
  • ๐Ÿ“Š Comprehensive EDA with visualization notebooks
  • ๐Ÿ”„ Batch and single-image prediction support
  • ๐Ÿ“ˆ Automatic metrics tracking with loss curves and performance reports
  • ๐Ÿ“š Auto-generated documentation using pdoc3
  • ๐ŸŽ“ LaTeX report generation for academic documentation

๐Ÿš€ Quick Start

# Install uv
curl -LsSf https://astral.sh/uv/install.sh | sh

# Activate uv
source $HOME/.local/bin/env

Interactive demo

# Create Virtual Environment
uv venv garbage-env --python 3.12

# Install PyPI package
uv pip install -U -i https://pypi.org/simple/ garbage-classifier

# Launch demo
garbage-app

Full project

# Clone the repository
git clone https://github.com/NeoLafuente/garbage_classifier.git
cd garbage_classifier

# Sync dependencies
uv sync

# Train the model
uv run source/train.py

# Make a prediction
uv run source/predict.py path/to/image.jpg

๐Ÿ“ฆ Dataset

The project uses the Garbage Classification Dataset from Kaggle.

Dataset Preparation

The notebook notebooks/create_sample_dataset.ipynb automatically prepares the dataset:

  • Downloads the Garbage Classification Dataset
  • Creates the sample_dataset folder inside data/raw
  • Reduces dataset size for lightweight experimentation

Dataset Structure:

  • 6 classes: cardboard, glass, metal, paper, plastic, trash
  • Location: data/raw/Garbage_Dataset_Classification/
  • Metadata: Class distributions and image statistics in metadata.csv

๐Ÿ’ป Usage

Training

Train the GarbageClassifier model (ResNet18 with PyTorch Lightning):

uv run source/train.py

Output:

  • Model checkpoint: models/weights/model_resnet18_garbage.ckpt
  • Loss curves: models/performance/loss_curves/
  • Training logs with metrics (accuracy, precision, recall, F1-score)

Configuration: Edit source/utils/config.py to customize:

  • Batch size
  • Learning rate
  • Number of epochs
  • Train/validation split ratio

Prediction

Load the trained model and classify new images.

๐Ÿ“ธ Single Image Prediction

# Predict a specific image
uv run source/predict.py img.jpg

# Use default image from config
uv run source/predict.py

๐Ÿ“ Batch Folder Prediction

Process all images in a folder:

uv run source/predict.py data/test_images/
uv run source/predict.py ../new_samples/

Features:

  • Auto-detects valid image files (.jpg, .jpeg, .png, .bmp, .gif, .tiff, .tif)
  • Progress indicators for batch processing
  • Summary table with all predictions

Example Output:

Predicting: cardboard_sample.jpg
Class: cardboard | Confidence: 98.5%

๐Ÿ—๏ธ Model Architecture

GarbageClassifier

  • Backbone: ResNet18 (pretrained on ImageNet)
  • Framework: PyTorch Lightning
  • Input: 224x224 RGB images
  • Output: 6-class probability distribution

Custom Components

  1. GarbageDataModule: PyTorch Lightning DataModule for efficient data loading
  2. LossCurveCallback: Custom callback for tracking and saving loss curves
  3. GarbageClassifier: Main model class with training/validation logic

๐Ÿ“š Documentation

Auto-generated Documentation

HTML documentation is auto-generated from source code docstrings using pdoc3.

View Documentation:

  • Open docs/index.html in your browser

Regenerate Documentation:

uv run scripts/generate_docs.py

๐Ÿ““ Notebooks

Interactive Jupyter notebooks for exploration and analysis:

Notebook Description
create_sample_dataset.ipynb Dataset preparation and sampling
dataset_exploration.ipynb EDA with class distribution and visualizations
performance_analysis.ipynb Model evaluation, confusion matrices, error analysis

๐Ÿ“„ Reports

LaTeX Report

Academic-style report with methodology and results:

  • Source: reports/main.tex
  • Compiled PDF: reports/compiled/
  • Figures:
    • reports/figures/EDA/ - Exploratory data analysis
    • reports/figures/performance/ - Model metrics and evaluation

โš™๏ธ Configuration

Central configuration in source/utils/config.py:

# Dataset configuration
CLASSES = ['cardboard', 'glass', 'metal', 'paper', 'plastic', 'trash']
DATA_PATH = 'data/raw/Garbage_Dataset_Classification'

# Model hyperparameters
BATCH_SIZE = 32
LEARNING_RATE = 0.001
NUM_EPOCHS = 50

# Data split
TRAIN_RATIO = 0.8
VAL_RATIO = 0.2

๐Ÿ“ Project Organization

.
โ”œโ”€โ”€ data
โ”‚   โ”œโ”€โ”€ interim              # Intermediate data transformations
โ”‚   โ”œโ”€โ”€ processed            # Final preprocessed data
โ”‚   โ””โ”€โ”€ raw                  # Original unprocessed datasets
โ”‚       โ””โ”€โ”€ Garbage_Dataset_Classification
โ”‚           โ”œโ”€โ”€ images       # Image files by class
โ”‚           โ””โ”€โ”€ metadata.csv # Dataset statistics
โ”œโ”€โ”€ docs                     # Auto-generated HTML documentation
โ”œโ”€โ”€ models
โ”‚   โ”œโ”€โ”€ performance          # Loss curves and metrics
โ”‚   โ”‚   โ””โ”€โ”€ loss_curves
โ”‚   โ””โ”€โ”€ weights              # Trained model checkpoints (.ckpt)
โ”œโ”€โ”€ notebooks                # Jupyter notebooks for EDA and analysis
โ”‚   โ”œโ”€โ”€ create_sample_dataset.ipynb
โ”‚   โ”œโ”€โ”€ dataset_exploration.ipynb
โ”‚   โ””โ”€โ”€ performance_analysis.ipynb
โ”œโ”€โ”€ reports                  # LaTeX reports and figures
โ”‚   โ”œโ”€โ”€ compiled             # PDF reports
โ”‚   โ”œโ”€โ”€ figures
โ”‚   โ”‚   โ”œโ”€โ”€ EDA              # Exploratory analysis plots
โ”‚   โ”‚   โ””โ”€โ”€ performance      # Model evaluation plots
โ”‚   โ””โ”€โ”€ main.tex             # Main report source
โ”œโ”€โ”€ scripts                  # Utility scripts
โ”‚   โ””โ”€โ”€ generate_docs.py     # Documentation generator
โ”œโ”€โ”€ source                   # Main source code
โ”‚   โ”œโ”€โ”€ predict.py           # Prediction script
โ”‚   โ”œโ”€โ”€ train.py             # Training script
โ”‚   โ””โ”€โ”€ utils
โ”‚       โ”œโ”€โ”€ config.py        # Configuration file
โ”‚       โ””โ”€โ”€ custom_classes   # Model implementations
โ”‚           โ”œโ”€โ”€ GarbageClassifier.py
โ”‚           โ”œโ”€โ”€ GarbageDataModule.py
โ”‚           โ””โ”€โ”€ LossCurveCallback.py
โ”œโ”€โ”€ pyproject.toml           # Project dependencies (uv)
โ””โ”€โ”€ README.md

Note: dummy.txt files are placeholders to preserve empty folder structure in Git.


๐Ÿค Contributing

Contributions are welcome! Please feel free to submit a Pull Request.

  1. Fork the repository
  2. Create your feature branch (git checkout -b feature/AmazingFeature)
  3. Commit your changes (git commit -m 'Add some AmazingFeature')
  4. Push to the branch (git push origin feature/AmazingFeature)
  5. Open a Pull Request

๐Ÿ“ License

This project is licensed under the MIT License - see the LICENSE file for details.


๐Ÿ™ Acknowledgments


๐Ÿ“ง Contact

Neo Lafuente - @NeoLafuente

Project Link: https://github.com/NeoLafuente/garbage_classifier


Made with โค๏ธ and โ™ป๏ธ for a cleaner planet

Project details


Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Source Distribution

garbage_classifier-1.0.0.tar.gz (57.5 kB view details)

Uploaded Source

Built Distribution

If you're not sure about the file name format, learn more about wheel file names.

garbage_classifier-1.0.0-py3-none-any.whl (54.9 kB view details)

Uploaded Python 3

File details

Details for the file garbage_classifier-1.0.0.tar.gz.

File metadata

  • Download URL: garbage_classifier-1.0.0.tar.gz
  • Upload date:
  • Size: 57.5 kB
  • Tags: Source
  • Uploaded using Trusted Publishing? No
  • Uploaded via: twine/6.2.0 CPython/3.11.13

File hashes

Hashes for garbage_classifier-1.0.0.tar.gz
Algorithm Hash digest
SHA256 a8a0a870403659ba7ba608e3f82a381c374bdaa547c39ff8e222907a261b3529
MD5 3a74eea8a585dab1c1a37d98dd7460fc
BLAKE2b-256 dc5d326c4f9c4e5271d28458c651353c95bb301c04138fef966d455f4e4c540a

See more details on using hashes here.

File details

Details for the file garbage_classifier-1.0.0-py3-none-any.whl.

File metadata

File hashes

Hashes for garbage_classifier-1.0.0-py3-none-any.whl
Algorithm Hash digest
SHA256 de27f6ea2dd3b4d809a2d6168439c904188d07c4e5cb97784bfe5bad3d132e27
MD5 a4e6cbf62fe89b73e856aff71365e0ac
BLAKE2b-256 746dd717641cef8e1589910ba90fcb7c231c64d268f12055c76fdfcd7e914a9e

See more details on using hashes here.

Supported by

AWS Cloud computing and Security Sponsor Datadog Monitoring Depot Continuous Integration Fastly CDN Google Download Analytics Pingdom Monitoring Sentry Error logging StatusPage Status page