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
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
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 Google Colab. You can find the final dataset splitted into train and test rootstrap-org/waste-classifier in a Google Drive folder.
Split into train/test You can find the code for the classifier resnet-model.ipynb load the model 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 If you are using the dataset, please give a citation of this repository. The dataset can be downloaded at rootstrap-org/waste-classifier.
In the second training, another set of images was collected using the TACO dataset 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 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.