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