Classifying Waste Images with Machine Learning

How AI can help us reduce, reuse, recycle

Photo by Jasmin Sessler on Unsplash

Table of Contents

1. Introduction
2. Data Source
3. Setup
4. Data Preprocessing
5. Training the Models
6. Results
7. Conclusion
8. References

Introduction

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.

Data Source

This is the dataset that I’ll be using:

Waste Classification data

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:

Imbalanced Data : How to handle Imbalanced Classification Problems

Setup

First things first, I import the python libraries that we’ll need throughout this project:

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import 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.

Data Preprocessing

To start, I’ll import the libraries we’ll need for preprocessing and building the machine learning models:

from tensorflow.keras import layers
from tensorflow.keras import Model
from tensorflow.keras.preprocessing.image import ImageDataGenerator
from 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:

ML Practicum: Image Classification | Machine Learning Practica

If you’re using Google Colab, I recommend using the GPU hardware accelerator to speed up the training process.

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

Advanced Guide to Inception v3 on Cloud TPU | Google Cloud

Results

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 evaluate():

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:

From my previous article on sarcasm detection:

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…

The same model performance evaluation can be done for our transfer learning model.

Conclusion

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:

Keras documentation: Keras Applications

References

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:

[1] Python Programming | Deep Learning basics with Python, TensorFlow and Keras by sentdex

[2] Analytics Vidhya | Everything you Should Know about Confusion Matrix for Machine Learning by Aniruddha Bhandari