Classifying Pokémon Images with Machine Learning

A convolutional neural network (CNN) walkthrough with code

Photo by Thimo Pedersen on Unsplash

Table of Contents

1. Introduction
2. Data Source
3. Setup
4. Exploratory Data Analysis
5. Data Preprocessing
6. Training the Model
7. Results
8. Conclusion
9. References

Introduction

As a machine learning beginner and Pokémon fan, what better way to explore data visualization tools and neural networks than Pokémon classification?

My goal is to build a classifier that can predict whether a Pokémon is a fire-type or a water-type based on its image.

To do this, I’ll use a convolutional neural network, a class of deep neural networks most commonly applied to analyzing visual imagery.

Data Source

I’ll be using this dataset from Kaggle:

Pokemon Image Dataset

It contains images of all Pokémon from generation 1 to 7, along with their primary and secondary types as a CSV file.

Setup

First things first, I import the necessary python libraries.

import numpy as np
import pandas as pd
import seaborn as sns
from matplotlib import pyplot as plt
import matplotlib.image as mpimg
import zipfile
import os
from pathlib import Path
import re
import tensorflow as tf
from tensorflow import keras
from keras import layers

Now I create a Pandas DataFrame with the file path and primary type of each Pokémon. I will not consider the secondary type to simplify the task. If you’re using Google Colab, I recommend uploading the dataset as a zip file and extracting it within the notebook.

Alternatively, you can remove the Pokémon that have a secondary type before dropping the Type2 column, as the secondary type may influence the appearance of Pokémon. From my trial runs, however, there are only 371 Pokémon that do not have a secondary type. The full dataset is already small for training a convolutional neural network, so further decreasing the number of data points may decrease the accuracy of the classifier despite having removed a potential source of noise.

Exploratory Data Analysis

To get a better sense of the dataset I’m working with, I’ll look at a few statistics and sample images.

I’ll start by creating a custom color palette based on the primary types of each Pokémon. You can find the source for Pokémon color templates here.

Next, I’ll plot the number of each primary type of Pokémon.

From this graph, we see that this dataset only has 3 flying-type Pokémon. With only 3 data points, it would be hard to train a model that can accurately identify flying-type Pokémon. On the other hand, the greatest number we have of any primary type is water-type with 105 Pokémon!

To see the variation between primary types, I’ll display one of each primary type of Pokémon:

For some Pokémon, it’s hard to see the connection between their appearance and their primary type. For others though, it’s much more obvious! For dark types, we can expect darker colors. For fire types, we can expect red or orange colors. But there are exceptions to these trends. To see the variation within a primary type, I’ll display images of the first 18 water types:

The general color trend for water types is blue as expected, but there are quite a few exceptions, like Goldeen in the top-left corner with its orange and white scales.

Data Preprocessing

Now that we have a better understanding of our data, it’s time to prepare the dataset for our machine learning model!

To start, I’ll limit the dataset to fire and water-type Pokémon:

Now I’ll generate training and validation sets of scaled images for feeding into our model. We use the training set to train our model and the validation set, a dataset that the model has never seen before, to evaluate the performance of our model later on. Here, I chose to split the dataset into training and validation sets in an 8:2 ratio. I seemed to get better results when I shuffled the dataset beforehand, but you can always shuffle after separating the data into training and validation sets.

From the output, we see we have 122 Pokémon in our training set and 30 Pokémon in our validation set.

Training the Model

Onto building and compiling our classifier:

Now we’re ready to fit our model! From the data-limiting step, we saw that there are more water-types than fire-types. This class imbalance is a problem, as the model may learn that the best thing to do is to only predict the more common class.

To prevent this, we will use AUC (Area under the ROC Curve) as a metric in addition to accuracy and loss. ROC curves are a popular metric for imbalanced binary classification problems as they are not biased towards the majority or minority class. You can read more about AUC and ROC curves here:

Classification: ROC Curve and AUC | Machine Learning Crash Course

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

Results

Now that we have trained our machine learning model, it’s time to visualize our results!

First, let’s graph the accuracy and loss of our model during training:

From the accuracy plot, both training and validation accuracy jump at around 19 epochs. As for the loss plot, the training loss drops early on at around epoch 1, while validation loss starts off quite low.

Next, let’s test our classifier on the validation set. Normally, we would have a test set that we don’t touch in addition to the training and validation sets, but I chose not to do that here as our dataset is small.

We first get our label predictions from the model. These label predictions come in the form of probabilities between 0 and 1. We then convert the probabilities into binary values. In our model, the label 0 corresponds to fire types, while 1 corresponds to water types. We then compare our label predictions to the actual labels to determine the number of misclassified Pokémon.

To visualize the classification results, I’ll output the images of the classified Pokémon along with correctness labels.

Interestingly, one of the Pokémon that our model misclassified is Volcanion, the only dual-type Fire/Water Pokémon.

Conclusion

We’ve successfully trained a classifier that can predict whether a Pokémon is a fire or water-type! There is definitely more to explore in the Pokémon classification world though. Here are some optimizations that can potentially improve the accuracy of our model:

• Add data augmentation
• Find other Pokémon image datasets and combine them with the one used here to increase the number of data samples
• Test out and/or combine different solutions to the class imbalance
• Use transfer learning as demonstrated in this tutorial:

ML Practicum: Image Classification | Machine Learning Practica

Besides improving accuracy, here are some ideas to build on this project:

• Determine both the primary and secondary type of a Pokémon based on its image
• Try different combinations of Pokémon to classify, or expand the problem to classify all types of Pokémon!

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] Kaggle | Movies and TV shows: EDA by Ruchi Bhatia

[2] Kaggle | Fire/Water Pokémon Classification by Gabriel Atkin

[3] Google Developers | ML Practicum: Image Classification