| --- |
| 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. |
