Image classification is a category of pattern recognition. It classifies images according to the relationship between the neighboring pixels. In other words, it uses contextual information to organize images and is quite popular among different technologies. It’s a prominent topic in Deep Learning, and if you’re learning about it, you’d surely enjoy this article.
Here, we’ll perform the TensorFlow image classification. We’ll build a model, train it, and then enhance its accuracy to classify images of cacti. TensorFlow is an open-source machine learning platform and a product of Google.
Let’s get started.
Table of Contents
Install TensorFlow 2.0
First, you’ll need to install TensorFlow on Google Colab. You can install it through pip:
!pip install tensorflow-gpu==2.0.0-alpha0
Then we’ll verify the installation:
import tensorflow as tf
# Output: 2.0.0-alpha0
After the verification, we can load the data by using the tf.data.dataset. We’ll build a classifier that determines whether an image contains a cactus or not. The cactus has to be columnar.We can use the Cactus Aerial Photos dataset for this purpose. Now, we’ll load the file paths along with their labels:
train_csv = pd.read_csv(‘data/train.csv’)
# Prepend image filenames in train/ with relative path
filenames = [‘train/’ + fname for fname in train_csv[‘id’].tolist()]
labels = train_csv[‘has_cactus’].tolist()
train_filenames, val_filenames, train_labels, val_labels =
Once we have the labels and filenames, we are ready to create the tf.data.Dataset objects:
train_data = tf.data.Dataset.from_tensor_slices(
val_data = tf.data.Dataset.from_tensor_slices(
At the moment, our dataset doesn’t have the actual images. It only has their filenames. We’ll need a function to load the necessary images and process them so we can perform TensorFlow image recognition on them.
IMAGE_SIZE = 96 # Minimum image size for use with MobileNetV2
BATCH_SIZE = 32
# Function to load and preprocess each image
def _parse_fn(filename, label):
img = tf.io.read_file(img)
img = tf.image.decode_jpeg(img)
img = (tf.cast(img, tf.float32)/127.5) – 1
img = tf.image.resize(img, (IMAGE_SIZE, IMAGE_SIZE))
return img, label
# Run _parse_fn over each example in train and val datasets
# Also shuffle and create batches
train_data = (train_data.map(_parse_fn)
val_data = (val_data.map(_parse_fn)
Building a Model
In this TensorFlow image classification example, we’ll create a transfer learning model. These models are fast as they can use existing image classification models that have undergone training before. They only have to retrain the upper layer of their network as this layer specifies the class of the required image.
We’ll use the Keras API of TensorFlow 2.0 to create our image classification model. And for the transfer learning purposes, we’ll use the MobileNetV2 to be the attribute detector. It’s the second version of MobileNet and is a product of Google. It’s lighter in weight than other models such as Inception and ResNet and can run on mobile devices. We’ll load this model on ImageNet, freeze the weights, add a classification head and run it without its top layer.
IMG_SHAPE = (IMAGE_SIZE, IMAGE_SIZE, 3)
# Pre-trained model with MobileNetV2
base_model = tf.keras.applications.MobileNetV2(
# Freeze the pre-trained model weights
base_model.trainable = False
# Trainable classification head
maxpool_layer = tf.keras.layers.GlobalMaxPooling2D()
prediction_layer = tf.keras.layers.Dense(1, activation=’sigmoid’)
# Layer classification head with feature detector
model = tf.keras.Sequential([
learning_rate = 0.0001
# Compile the model
You should use TensorFlow optimizers if you are going to train tf.keras models. The optimizers in tf.keras.optimizers and tf.train APIs are together in TensorFlow 2.0’s tf.keras.optimizers. In TensorFlow 2.0, many of the original optimizers of tf.keras have received upgrades and replacements for better performance. They enable us to apply optimizers without compromising with the performance and save time as well.
Learn data science courses from the World’s top Universities. Earn Executive PG Programs, Advanced Certificate Programs, or Masters Programs to fast-track your career.
Training the Model
After we’ve built the model, we can teach it. The tf.keras API of TensorFlow 2.0 supports the tf.data API, so you have to use the tf.data.Dataset objects for this purpose. It can perform the training efficiently, and we wouldn’t have to make any compromises with the performance.
num_epochs = 30
steps_per_epoch = round(num_train)//BATCH_SIZE
val_steps = 20
steps_per_epoch = steps_per_epoch,
After 30 epochs, the model’s accuracy increases substantially, but we can improve it further. Remember, we mentioned freezing the weights during transfer learning? Well, now that we have trained the classification head, we can unfreeze those layers and fine-tune our dataset further:
# Unfreeze all layers of MobileNetV2
base_model.trainable = True
# Refreeze layers until the layers we want to fine-tune
for layer in base_model.layers[:100]:
layer.trainable = False
# Use a lower learning rate
lr_finetune = learning_rate / 10
# Recompile the model
optimizer = tf.keras.optimizers.Adam(lr=lr_finetune),
# Increase training epochs for fine-tuning
fine_tune_epochs = 30
total_epochs = num_epochs + fine_tune_epochs
# Fine-tune model
# Note: Set initial_epoch to begin training after epoch 30 since we
# previously trained for 30 epochs.
steps_per_epoch = steps_per_epoch,
initial_epoch = num_epochs,
30 epochs later, the model’s accuracy improves further. WIth more epochs, we saw more improvement in the accuracy of the model. Now, we have a proper TensorFlow image recognition model that can recognize columnar cacti in images with a high accuracy.
Also Read: Tensorflow Project Ideas for Beginners
Learn More About TensorFlow Image Classification
The highly functional APIs of TensorFlow and its capabilities make it a powerful technology for any programmer to wield. Its high-level APIs also remove its general complexity, making it easier to use.
Are you interested in learning more about TensorFlow, image classification, and related topics? Then we recommend you to check out IIIT-B & upGrad’s PG Diploma in Machine Learning & AI which is designed for working professionals and offers 450+ hours of rigorous training, 30+ case studies & assignments, IIIT-B Alumni status, 5+ practical hands-on capstone projects & job assistance with top firms.