|
--- |
|
datasets: |
|
- rootstrap-org/waste-classifier |
|
library_name: fastai |
|
tags: |
|
- classification |
|
- waste |
|
- fastai |
|
- resnet |
|
license: mit |
|
--- |
|
|
|
--- |
|
|
|
### Model Description |
|
A machine learning model for waste classification |
|
|
|
- **Developed by:** rootstrap |
|
- **Model type:** classifier |
|
- **License:** mit |
|
|
|
## Waste Classifier Model |
|
The aim is to build a model for waste classification that identifies among the different classes: |
|
|
|
- cardboard |
|
- compost |
|
- glass |
|
- metal |
|
- paper |
|
- plastic |
|
- trash |
|
|
|
This machine learning model will help people to improve their decision when classifying trash |
|
|
|
|
|
The model was built using **Fastai** is a deep learning library which provides high-level components that can quickly and easily provide state-of-the-art results in standard deep learning domains. |
|
|
|
It has two main design goals: |
|
|
|
To be approachable and rapidly productive |
|
To be also configurable. |
|
|
|
### Model Sources |
|
|
|
- **Repository:** [https://github.com/rootstrap/fastai-waste-classifier](https://github.com/rootstrap/fastai-waste-classifier) |
|
|
|
## Uses |
|
|
|
At Rootstrap we classify waste. We found that people were struggled to classify correctly, |
|
and then we end up not recycling most of the generated waste at the office, since if there were items in the wrong basket, |
|
all the basket should not be classified. |
|
Because of this reason, we created an app to help people at our company to classify waste. |
|
|
|
### Direct Use |
|
|
|
```bash |
|
model = load_learner("result-resnet50.pkl") |
|
model.predict() |
|
``` |
|
|
|
## Bias, Risks, and Limitations |
|
|
|
Take into account that the model has been trained firstly trained with images that mostly have white background that contains only one object. In a further training the model was trained with a more varied range of images, which improved the bias a little bit. But the model still is incapable of detecting more than one object per image. So, |
|
you won't be able to classify several objects at the same time, or pictures that have one object above another. |
|
This does not aim to replace humans to classify trash, at least for now, instead it is a tool to help them to better classify. |
|
|
|
## Training Details |
|
|
|
### Training Data |
|
The data is already splitted in train and test folders. Inside each folder contains one folder for each class. Those images were obtained using Bing searcher using the api HTTP. |
|
You can find the code used to download the images at [this](https://colab.research.google.com/drive/1JvAYFx1DIEi1MMyI-tuCfE2eHMSKisKT?usp=sharing) Google Colab. |
|
You can find the final dataset splitted into train and test [rootstrap-org/waste-classifier](https://huggingface.co/datasets/rootstrap-org/waste-classifier) in a Google Drive folder. |
|
|
|
Split into train/test |
|
You can find the code for the classifier [resnet-model.ipynb](https://github.com/rootstrap/fastai-waste-classifier/blob/main/resnet-model.ipynb) |
|
load the model [result-resnet50.pkl](https://huggingface.co/rootstrap-org/waste-classifier/blob/main/result-resnet50.pkl) |
|
|
|
The model has been trained to classify waste into 6 classes. |
|
This repository contains the dataset that we collected. |
|
The dataset used for the first training consisted of 2527 images: |
|
- 501 glass |
|
- 594 paper |
|
- 403 cardboard |
|
- 482 plastic |
|
- 410 metal |
|
- 137 trash |
|
|
|
The dataset was built joining images search on the internet using Bing plus some images taken manually, in addition to images from [Trashnet dataset](https://github.com/garythung/trashnet) |
|
If you are using the dataset, please give a citation of this repository. The dataset can be downloaded at [rootstrap-org/waste-classifier](https://huggingface.co/datasets/rootstrap-org/waste-classifier). |
|
|
|
In the second training, another set of images was collected using the [TACO dataset](http://tacodataset.org/) as a starting point, along with other datasets in the web. |
|
A selection of proper training images was made from those datasets to build the training dataset. |
|
|
|
### Training Procedure |
|
You can find the code for training at [resnet-model.ipynb](https://github.com/rootstrap/fastai-waste-classifier/blob/main/resnet-model.ipynb) |
|
Training the model by fine tuning resnet50, a convolutional neural network that has 50 layers. It is pretrained on ImageNet dataset and can classify images into 1000 object categories. |
|
The expected input size is 224x224. |
|
|
|
## Evaluation and Results |
|
After the first training, the model reached an accuracy of 0.98. |
|
However, that first model had a lot of bias and limitations. |
|
The final model trained with more images, many of them from the TACO dataset, reached an accuracy of 89%. |
|
That model still has some limitation, but it's a lot more flexible than the first one. |
