Custom Vision: Chapter 3 - Image Classification - Tensorflow

Computer Vision Jul 26, 2021

In our previous posts we explored the Roadmap of Custom Vision and covered some of the Pitstops on our way to this Chapter.

So What is this Chapter about? Well we know the importance of Custom Vision to the Industry. We are also aware of the various hurdles one has to cross in order to be able to create a finished product.

As promised in this series we will cover all the necessary steps that you may tackle in your journey of building a Custom Vision Solution.

Today's article is about using one of the leading Open Source Frameworks - Tensorflow to build a Custom Neural Network Architecture for Image Classification.

There are thousands of posts floating around on the internet guiding us to re-use the pre-trained models to fine tune it for our needs.

But what if you were aspiring to be a Deep Learning Engineer. You would potentially have to build Novel Models and understand the nuances involved with the same.

To embark on this long journey, today we will take the first step to identify images using Custom Neural Architecture.

Let's get Started ...

Step 0: Install all the necessary Packages

pip3 install tensorflow
pip3 install numpy

Step 1: Import all the necessary Packages

import os
import random
import numpy as np

from tensorflow.keras.preprocessing.image import load_img
from tensorflow.keras.layers import Input, Conv2D, Flatten, Dense
from tensorflow.keras.models import Model

As you can see, the primary framework that we are using to Load Images, to create Neural Network and to Define the Model is Tensorflow.

We are importing several Layers from Tensorflow which we will be using to build our Custom Architecture. We will get into further details in the upcoming steps.

Step 2: Download the Dataset and Relax

For today's article we will be using Intel Image Classification Dataset.

Download the Dataset and store it on your local system. Be sure to remember the path where you store the Dataset as we will be using it for our training job. Sit back and Relax as this might take some time depending on your internet speed.

Let's talk about the Dataset a bit. It has over 14000 labelled images, tagged and placed in their respective folders. It includes Images from the following categories -

  1. Forest
  2. Buildings
  3. Glacier
  4. Street
  5. Mountain
  6. Sea

Once you have downloaded the Images, pass the PATH where you have stored the Dataset.

dataset_path = "/Users/vsatpathy/Desktop/off_POCs/intel-image-classification/seg_train/"

Step 3: Read and Store the Meta Data needed

As all the data is sitting on your local system, we can just store the Labels and the corresponding file paths separately, which can be used to train the model.

As and when needed, we can load the images and pass it to the model. This way you wouldn't have to refer to the complete Image File at all times.

def data_reading(dataset_path):
    all_labels = {}
    all_images = []
    for folders in os.listdir(dataset_path):
        if folders != ".DS_Store":
            folder_path = os.path.join(dataset_path, folders)
            all_labels[folders] = len(all_labels)
            for images in os.listdir(folder_path):
                if images != ".DS_Store":
                    image_path = os.path.join(folder_path, images)

    rev_labels = {}
    for key, labels in all_labels.items():
        rev_labels[labels] = key

    return all_images, all_labels, rev_labels

The above function return the following -

  1. all_images - A list of all the paths of where the images are stored
  2. all_labels - A Dictionary of all the labels with their corresponding token value
  3. rev_labels - A Dictionary where the Key and Values are reversed as all_labels. This is used for inferencing, if needed in the future.

Step 4: Create a Data Generator

In general if you want to speed up your development process, Keras provides its own Data Generator. But as we are firm that we want to build everything Custom, here is an example of how you can create your own Data Generators.

def data_loader(bs, data, y_lab, image_input_shape):
    while True:
        images = []
        labels = []
        while len(images) < bs:
            indice = random.randint(0, len(data) - 1)
            target = data[indice].split("/")[-2]

            test_img = np.asarray(load_img(data[indice], target_size=image_input_shape))
            img = np.divide(test_img, 255.0)

        yield np.asarray(images), np.asarray(labels)

Data Generators are very crucial when you are dealing with massive volumes of Data.

If you fail to use Generator and instead choose to load the complete Dataset in your system memory, it may lead to an overflow and your Training would fail.

The highlight of Data Generators are, they use the YIELD function of Python to only store those data which are a part of the Batch which is being processed, thereby freeing up a lot of memory.

A Data Generator returns the following -

  1. Input Data Array for the Model
  2. Ground Truth Values or Labels in this case for the Model

Step 5: Create the Model

Now that we have processed the Data, it's time for the Neural Expert to rise and shine.

We create a simple Convolutional Neural Network (CNN) for our implementation, using the above mentioned Layers.

def model_arc(y_labels, image_inp_shape):
    inp_layer_images = Input(shape=image_inp_shape)

    conv_layer = Conv2D(filters=64, kernel_size=(2, 2), activation="relu")(
    flatten_layer = Flatten()(conv_layer)
    out_layer = Dense(len(y_labels), activation="softmax")(flatten_layer)
    model = Model(inp_layer_images, out_layer)
        optimizer="adam", loss="sparse_categorical_crossentropy", metrics=["accuracy"]
    return model

As you can see, we are not just declaring the model in the previous function but also compiling it with the necessary Loss Functions and Optimisers.

After running the above function, you would notice the Model Summary being printed on your terminal.

Step 6: Create your Train Function

All that's left to be done is create an entry point function for your code.

def train(dataset_path, batch_size, epochs, input_shape):
    all_images, all_labels, rev_labels = data_reading(dataset_path=dataset_path)

    print("target_encodings: ", all_labels)
    print("Number of training images: ", len(all_images))

    train_generator = data_loader(
        bs=batch_size, y_lab=all_labels, image_input_shape=input_shape, data=all_images

    model = model_arc(y_labels=all_labels, image_inp_shape=input_shape)

        steps_per_epoch=(len(all_images) // batch_size),

The job of this function is to call the respective steps in the chronological order as expected for Training the Model.

Step 6: Let it Train

batch_size = 8
epochs = 10
input_shape = (100, 100, 3)

Epoch 1/2
1754/1754 [==============================] - 62s 35ms/step - loss: 1.0267 - accuracy: 0.7274
Epoch 2/2
1754/1754 [==============================] - 60s 34ms/step - loss: 0.4079 - accuracy: 0.8796

You can set the Hyper-parameters to your liking. As a recommendation do try out different parameters to see how it effects the training results and Accuracy.

Enjoy your State of the Art Accuracy with your Custom Model for your Personalised Dataset.

Congratulations ...

If you have followed the above steps then you should have a Custom trained model on your hand.

I hope this article finds you well and was easy to follow. In our further posts we will be looking at how we can create Custom Models using Cloud Native Solutions for working more closely with Industry Requirements.

STAY TUNED for more such content 😁


Vaibhav Satpathy

AI Enthusiast and Explorer

Great! You've successfully subscribed.
Great! Next, complete checkout for full access.
Welcome back! You've successfully signed in.
Success! Your account is fully activated, you now have access to all content.