The Waste Detection System is a computer vision-based application that identifies and classifies different types of waste materials, including cardboard, glass, metal, paper, and plastic. This system utilizes PyTorch and OpenCV to create a waste detection model, preprocess data, train the model, and perform real-time waste classification from a live video feed.
Model Testing with Recorded Data
- Clone repo
git clone https://github.com/Micz26/CNN-Waste-Detector.git
- Create venv and install dependencies
poetry install
- Install pre-commit
pre-commit install
- Activate venv
poetry shell
- [Optional] install dependecies for working with jupyter notebooks
poetry install --extras "jupyter"
- [Optional] install dependecies for lint checking
poetry install --extras "lint"
After Installation, you can run scripts below. Specify arguments in <>.
- Script for training my models (
BaseModel
orModelRegularized
)
python scripts/train_model.py <model> <epochs> <on_save> <save_path>
- Script for fine-tuning vgg16
python scripts/fine_tune_vgg16.py <epochs> <on_save> <save_path>
- Script for recording real-time predictions
python scripts/record.py <model> <model_path>
- Script for oversamplnig empty images
python scripts/add_empty_images_to_data.py
The project is divided into several key components, including eda, data preprocessing, model training. Additionaly i fined-tuned vgg16
model on the data and presented it's perforemance in real-time waste classification using a webcam.
The data used in this project was obtained from the TrashNet Dataset. I enhanced the dataset by adding images labeled as 'nothing,' by taking few pictures of my desk and oversampling them. To ensure the model effectively learns to identify when there isn't any trash present.
I developed two models, BaseModel
and ModelRegularized
. Both models consist of six convolutional layers, each followed by ReLU activation, batch normalization, and max pooling, whcih are then followed by linear transformations. The primary difference between the two models is that ModelRegularized includes dropout layers to reduce overfitting. Additionally, I fine-tuned a pre-trained VGG16
model on the dataset.
The BaseModel
achieved an accuracy of 75%, while ModelRegularized
reached 70%. VGG16
significantly outperformed my models, achieving approximately 96% accuracy.
This part of the project demonstrates real-time waste classification using a webcam. You can watch it on youtube by clicking on the video:
Model Testing with Recorded Data
The data used in this project was obtained from the TrashNet Dataset, created by Gary Thung and Mindy Yang.
Please refer to the original repository for more information about the dataset and its creators.
This project is licensed under the MIT License - see the LICENSE file for details.