Finetunned ResNet18 for garbage image classification
Project description
๐๏ธ Garbage Classifier
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_datasetfolder insidedata/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
- GarbageDataModule: PyTorch Lightning DataModule for efficient data loading
- LossCurveCallback: Custom callback for tracking and saving loss curves
- 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.htmlin 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 analysisreports/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.
- Fork the repository
- Create your feature branch (
git checkout -b feature/AmazingFeature) - Commit your changes (
git commit -m 'Add some AmazingFeature') - Push to the branch (
git push origin feature/AmazingFeature) - Open a Pull Request
๐ License
This project is licensed under the MIT License - see the LICENSE file for details.
๐ Acknowledgments
- Dataset: Garbage Classification Dataset by zlatan599 on Kaggle
- Framework: PyTorch Lightning for the excellent deep learning framework
- Model: ResNet18 architecture from torchvision
๐ง Contact
Neo Lafuente - @NeoLafuente
Project Link: https://github.com/NeoLafuente/garbage_classifier
Project details
Release history Release notifications | RSS feed
Download files
Download the file for your platform. If you're not sure which to choose, learn more about installing packages.
Source Distribution
Built Distribution
Filter files by name, interpreter, ABI, and platform.
If you're not sure about the file name format, learn more about wheel file names.
Copy a direct link to the current filters
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
| Algorithm | Hash digest | |
|---|---|---|
| SHA256 |
a8a0a870403659ba7ba608e3f82a381c374bdaa547c39ff8e222907a261b3529
|
|
| MD5 |
3a74eea8a585dab1c1a37d98dd7460fc
|
|
| BLAKE2b-256 |
dc5d326c4f9c4e5271d28458c651353c95bb301c04138fef966d455f4e4c540a
|
File details
Details for the file garbage_classifier-1.0.0-py3-none-any.whl.
File metadata
- Download URL: garbage_classifier-1.0.0-py3-none-any.whl
- Upload date:
- Size: 54.9 kB
- Tags: Python 3
- Uploaded using Trusted Publishing? No
- Uploaded via: twine/6.2.0 CPython/3.11.13
File hashes
| Algorithm | Hash digest | |
|---|---|---|
| SHA256 |
de27f6ea2dd3b4d809a2d6168439c904188d07c4e5cb97784bfe5bad3d132e27
|
|
| MD5 |
a4e6cbf62fe89b73e856aff71365e0ac
|
|
| BLAKE2b-256 |
746dd717641cef8e1589910ba90fcb7c231c64d268f12055c76fdfcd7e914a9e
|