Skip to content

A Pytorch project for garbage classification using the EfficientNet-B6 model to achive a 95.78% accuracy on the test set. 😊

License

Notifications You must be signed in to change notification settings

Aydinhamedi/Pytorch-Garbage-Classification-V2

Repository files navigation

Garbage Classification V2 with PyTorch

License: MIT Ruff

A Pytorch project for garbage classification using the EfficientNet-B6 model to achieve a 95.78% accuracy on the test set.

Important

This project is a new version of the original project, which can be found here but with a significantly improved training process + code and a different dataset.

πŸ˜‰ Bonus

This project is not hard coded for this specific dataset, so it can be used for any image classification task and it has all the necessary tools to train a model from scratch and more. (I will make a pytorch classification template soon)

πŸ“¦ Release

Newest release πŸ“ƒ

πŸ“‚ Dataset

The dataset used for this project is the Garbage Classification from Kaggle. It contains images of garbage, divided into six categories.

Data Structure

β”œβ”€β”€β”€Database
β”‚   └───Data # Put all the folders with images here
#       Example ⬎
β”‚       β”œβ”€β”€β”€battery
β”‚       β”œβ”€β”€β”€biological
β”‚       β”œβ”€β”€β”€brown-glass
β”‚       ...
β”‚       └───white-glass

πŸ§ͺ Model

I used the EfficientNet-B6 model for this project. EfficientNet-B6 is a convolutional neural network that is pretrained on the ImageNet dataset. It is known for its efficiency and high performance on a variety of image classification tasks. Original paper

πŸ”° Installation

To run the code in this repository, you will need to install the required libraries. You can do this by running the following command:

pip install -r requirements.txt

Warning

The requirements are auto generated by pipreqs and may not contain all the necessary dependencies. like hidden ones like Tensorboard.

πŸš€ Usage

The main code for this project is in a Jupyter notebook named Main.ipynb. To run the notebook, use the following command:

jupyter notebook Main.ipynb

πŸ“ƒ Results

Metric Value
Loss 0.0330466
F1 Score (macro) 0.95472
Precision (macro) 0.952111
Recall (macro) 0.957959
AUROC 0.993324
Accuracy 0.957839
Cohen's Kappa 0.948292
Matthews Correlation Coefficient 0.948374

alt text alt text

πŸ“š License

 Copyright (c) 2024 Aydin Hamedi
 
 This software is released under the MIT License.
 https://opensource.org/licenses/MIT