Image Classification with TensorFlow in Machine Learning

In this article, I am going to explain how we can train a neural network model for the task of image classification with TensorFlow. For those new to TensorFlow, TensorFlow is an end-to-end open-source platform for machine learning. It has a comprehensive and flexible ecosystem of tools, libraries, and community resources that allow researchers to push cutting-edge advancements in ML, and developers to easily build and deploy machine learning-based applications.

What is Image Classification?

Image classification is the process of categorizing and labelling groups of pixels or vectors in an image according to specific rules. The categorization law can be designed using one or more spectral or textural characteristics.

Also, Read – Why Python is Better than R?

Image Classification with TensorFlow

Now, Image Classification can also be done by using less complex models provided by Scikit-Learn, so why TensorFlow. By using TensorFlow we can build a neural network for the task of Image Classification. By building a neural network we can discover more hidden patterns than just classification. Now let’s get started with the task of Image Classification with TensorFlow by importing some necessary packages:

# TensorFlow and tf.keras
import tensorflow as tf
from tensorflow import keras

# Helper libraries
import numpy as np
import matplotlib.pyplot as plt
Code language: PHP (php)

Import the Fashion MNIST dataset

Fashion MNIST is intended as a drop-in replacement for the classic MNIST dataset—often used as the “Hello, World” of machine learning programs for computer vision. The MNIST dataset contains images of handwritten digits (0, 1, 2, etc.) in a format identical to that of the images of clothing that I will use for the task of image classification with TensorFlow.

The Fashion MNIST Dataset is an advanced version of the traditional MNIST dataset which is very much used as the “Hello, World” of machine learning. The MNIST dataset contains images of handwritten numbers (0, 1, 2, etc.) in the same format as the clothing images I will be using for the image classification task with TensorFlow. Now let’s import the Fashion MNIST dataset to get started with the task:

fashion_mnist = keras.datasets.fashion_mnist
(train_images, train_labels), (test_images, test_labels) = fashion_mnist.load_data()

In the dataset, each image is mapped into a single label. Since the class names are not defined in the dataset, we need to store them here so that we can use them later when viewing the images:

class_names = ['T-shirt/top', 'Trouser', 'Pullover', 'Dress', 'Coat',
               'Sandal', 'Shirt', 'Sneaker', 'Bag', 'Ankle boot']Code language: JavaScript (javascript)

Preprocess The Data

For this task of Image Classification with TensorFlow, the data must be preprocessed before training the neural network. If you inspect the first frame of the training set, you will find that the pixel values ​​are between 0 and 255:

plt.grid(False) language: CSS (css)

Now I’m going to scale these values ​​to a range of 0 to 1 before passing them to the neural network model. To do this, we need to divide the values ​​by 255. The training set and the test set should be preprocessed in the same way:

train_images = train_images / 255.0
test_images = test_images / 255.0

To verify that the data is in the correct format and to verify that we are ready to create and train the neural network for image classification with TensorFlow, let’s display the first 25 images of the training set and display the name of the class under each image:

for i in range(25):
    plt.xlabel(class_names[train_labels[i]]) language: PHP (php)

Image Classification with TensorFlow: Building Model

Now to Build the neural network for the task of Image Classification with TensorFlow, we first need to configure the model layers and then move forward with compiling the model.

Setting Up Layers

The basic building block of neural networks is its layers. Layers work by extracting the representations from data fed into them. Most of the deep learning, Models involves doing simple layers together. Now, let’s create the layers of our neural network:

model = keras.Sequential([
    keras.layers.Flatten(input_shape=(28, 28)),
    keras.layers.Dense(128, activation='relu'),
])Code language: JavaScript (javascript)

Compiling The Model

Now, let’s move forward with compiling our model:

              metrics=['accuracy'])Code language: PHP (php)

Image Classification with TensorFlow: Training Model

Now, let’s train the Neural Network for the task of Image Classification with TensorFlow, and make predictions on it:

#Fitting the Model, train_labels, epochs=10)
#Evaluating Accuracy
test_loss, test_acc = model.evaluate(test_images,  test_labels, verbose=2)
print('\nTest accuracy:', test_acc)Code language: PHP (php)

Test accuracy: 0.8817999958992004

#Make Predictions
probability_model = tf.keras.Sequential([model, 
predictions = probability_model.predict(test_images)
predictions[0]Code language: PHP (php)
array([1.1349098e-09, 1.0395625e-09, 3.4154518e-10, 8.3033120e-12,
       6.5739442e-10, 5.9645530e-03, 9.4151291e-09, 1.1747092e-02,
       8.7000714e-08, 9.8228824e-01], dtype=float32)

A prediction is an array of 10 numbers. They represent the “confidence” of the model that the image matches each of the 10 different garments. Let’s see which label has the highest confidence value:

np.argmax(predictions[0])Code language: CSS (css)
Output: 9

Thus, the model is most convinced that this image is an ankle boot, or class_names [9]. Examination of the test label shows that this classification is correct:

test_labels[0]Code language: CSS (css)

Now, I will create a helper function to plot our predictions:

def plot_image(i, predictions_array, true_label, img):
  true_label, img = true_label[i], img[i]


  predicted_label = np.argmax(predictions_array)
  if predicted_label == true_label:
    color = 'blue'
    color = 'red'

  plt.xlabel("{} {:2.0f}% ({})".format(class_names[predicted_label],

def plot_value_array(i, predictions_array, true_label):
  true_label = true_label[i]
  thisplot =, predictions_array, color="#777777")
  plt.ylim([0, 1])
  predicted_label = np.argmax(predictions_array)

  thisplot[true_label].set_color('blue')Code language: PHP (php)

Verify Predictions

Let’s look at the 0th frame of the predictions and the prediction table. The correct prediction labels are blue and the incorrect prediction labels are red:

i = 0
plot_image(i, predictions[i], test_labels, test_images)
plot_value_array(i, predictions[i],  test_labels)
image classification
# Plot the first X test images, their predicted labels, and the true labels.
# Color correct predictions in blue and incorrect predictions in red.
num_rows = 5
num_cols = 3
num_images = num_rows*num_cols
plt.figure(figsize=(2*2*num_cols, 2*num_rows))
for i in range(num_images):
  plt.subplot(num_rows, 2*num_cols, 2*i+1)
  plot_image(i, predictions[i], test_labels, test_images)
  plt.subplot(num_rows, 2*num_cols, 2*i+2)
  plot_value_array(i, predictions[i], test_labels)
plt.tight_layout() language: PHP (php)
image classification with tensorflow

Also, Read – Structured and Unstructured Data in Machine Learning.

The output looks great, only the boots are recognized wrong as sandals. I hope you liked this article on Image Classification with Machine Learning. Feel free to ask your valuable questions in the comments section below. You can also follow me on Medium to learn every topic of Machine Learning.

Follow Us:

Aman Kharwal
Aman Kharwal

I'm a writer and data scientist on a mission to educate others about the incredible power of data📈.

Articles: 1500

Leave a Reply