With the amount of waste produced daily worldwide, waste management is a massive problem. A significant part of this issue is waste classification. But what if we could use AI to automate the process?
My goal is to build machine learning models that can classify waste as organic or recyclable from images.
Similar to my past article on Pokémon classification, I’ll do this using a convolutional neural network.
This is the dataset that I’ll be using:
Waste Classification data
This dataset contains 22500 images of organic and recyclable objects
It is split into test and train directories that are both further divided into organic and recyclable folders.
The original dataset has more organic than recyclable images for both the training and validation sets. If we don’t remedy this class imbalance, the model may learn that it is best to predict only the majority class.
To prevent this, we’ll trim the organic data down to the size of the recyclable data. I did this within file explorer with the files sorted according to descending file size for randomness. Alternative solutions include:
- Using AUC (Area under the ROC Curve) as a metric in addition to accuracy and loss
- Random oversampling of the minority class
- Using class weights when training the model
You can read about more methods for dealing with imbalanced data here:
First things first, I import the python libraries that we’ll need throughout this project:
import numpy as npimport pandas as pdimport matplotlib.pyplot as pltimport matplotlib.image as mpimg
Then I set up the paths to each directory. If you’re using Google Colab, I recommend uploading the dataset as a zip file and extracting it within the notebook.
From the output, we have 9999 training and 1112 validation images for each class.
We’ll need to know what the shapes of the images are for our machine learning model, so let’s take a look at a few samples:
All images are RGB, but their dimensions vary, which needs to be accounted for when preprocessing the data. I’ve also displayed sample images from each class. The first four are from the organic directory, while the bottom four are recyclable.
To start, I’ll import the libraries we’ll need for preprocessing and building the machine learning models:
from tensorflow.keras import layersfrom tensorflow.keras import Modelfrom tensorflow.keras.preprocessing.image import ImageDataGeneratorfrom tensorflow.keras.optimizers import RMSprop
Now we’ll create our image data generators. We’ll rescale the data between 0 and 1 while resizing the images to be 150 by 150. From my trial runs, data augmentation did not improve model performance.
Training the Models
Now we’re ready to train our classifiers! I’ll be creating two machine learning models: a CNN from scratch and a transfer learning model. I based my models on this image classification practicum by Google Developers:
If you’re using Google Colab, I recommend using the GPU hardware accelerator to speed up the training process.
- CNN from Scratch
2. Transfer Learning Model
For my transfer learning model, I’ll make use of Google’s Inception v3 model. You can read more about Inception v3 here:
The CNN from scratch achieved slightly higher accuracy than the transfer learning model, so let’s visualize and evaluate its performance!
First, let’s test our classifier on the validation set:
Organic samples correspond to the label 0, while recyclable samples correspond to 1. Out of the 2224 validation images, our model misclassified 197.
To better visualize these results, let’s output the images of the classified waste along with correctness labels.
We can also evaluate our model’s performance using other metrics. Using
Our model’s accuracy on the validation set is about 91%, while the loss is about 0.25.
Using the scikit-learn library, we can also create a confusion matrix and a classification report with more details:
In a confusion matrix, entry i, j is the number of observations in group i but predicted to be in group j.
As for the classification report:
- Precision: the ability of the classifier to not label a negative sample as positive
- Recall: the ability of the classifier to find all positive samples
- f1-score: weighted average of precision and recall
- Support: number of samples classified in the given class
You can read more about these classification metrics here:
3.3. Metrics and scoring: quantifying the quality of predictions - scikit-learn 0.24.2…
Model selection and evaluation using tools, such as model_selection.GridSearchCV andmodel_selection.cross_val_score …
The same model performance evaluation can be done for our transfer learning model.
We’ve successfully built a classifier that can differentiate between images of organic waste and images of recyclable waste! But we still have a long way to go before solving the world’s waste crisis. Here are some ideas for taking this project further:
- Combine the dataset used here with other waste classification datasets to increase accuracy and/or the number of waste classes classified
- Test out and/or combine different solutions to the class imbalance
- Experiment with other transfer learning models! Here are the ones available with Keras:
In addition to the ones linked throughout this article, I wouldn’t have been able to complete this project without the help of these awesome examples and tutorials:
 Python Programming | Deep Learning basics with Python, TensorFlow and Keras by sentdex
 Analytics Vidhya | Everything you Should Know about Confusion Matrix for Machine Learning by Aniruddha Bhandari