How can Tensorflow be used to add a batch dimension and pass the image to the model using Python?

TensorFlow can be used to add a batch dimension and pass the image to the model by converting the image to a NumPy array and using np.newaxis to expand dimensions.

Read More: What is TensorFlow and how Keras work with TensorFlow to create Neural Networks?

A neural network that contains at least one layer is known as a convolutional layer. We can use the Convolutional Neural Network to build learning model.

Understanding Batch Dimensions

Most TensorFlow models expect input data with a batch dimension as the first axis. Even for single images, we need to add this batch dimension. The batch dimension allows the model to process multiple images simultaneously during training or inference.

Adding Batch Dimension to Images

Here's how to prepare an image and add the required batch dimension ?

import numpy as np
import tensorflow as tf

# Simulate loading an image (normally you'd use tf.keras.preprocessing.image.load_img)
# Creating a sample image array of shape (224, 224, 3)
grace_hopper = np.random.rand(224, 224, 3)

# Normalize the image to [0, 1] range
grace_hopper = np.array(grace_hopper) / 255.0
print("The dimensions of the image are")
print(grace_hopper.shape)

# Add batch dimension using np.newaxis
batched_image = grace_hopper[np.newaxis, ...]
print("The dimensions after adding batch dimension are")
print(batched_image.shape)

The dimensions of the image are
(224, 224, 3)
The dimensions after adding batch dimension are
(1, 224, 224, 3)

Passing Image to Model for Prediction

Once the batch dimension is added, you can pass the image to a pre-trained model ?

import numpy as np
import tensorflow as tf

# Create a simple mock classifier for demonstration
# In practice, you'd load a pre-trained model from TensorFlow Hub
class MockClassifier:
    def predict(self, x):
        # Simulate prediction returning logits for 1001 classes
        batch_size = x.shape[0]
        return np.random.rand(batch_size, 1001)

classifier = MockClassifier()

# Simulate an image
grace_hopper = np.random.rand(224, 224, 3)
grace_hopper = np.array(grace_hopper) / 255.0

print("The dimensions of the image are")
print(grace_hopper.shape)

# Add batch dimension and make prediction
result = classifier.predict(grace_hopper[np.newaxis, ...])
print("The dimensions of the resultant image are")
print(result.shape)

# Get the predicted class
predicted_class = np.argmax(result[0], axis=-1)
print("The predicted class is")
print(predicted_class)

The dimensions of the image are
(224, 224, 3)
The dimensions of the resultant image are
(1, 1001)
The predicted class is
543

Alternative Methods for Adding Batch Dimension

Key Points

• A batch dimension is added as the first axis using np.newaxis
• The image shape changes from (height, width, channels) to (1, height, width, channels)
• The model returns predictions with shape (batch_size, num_classes)
• Use np.argmax() to get the predicted class index from logits
• Image normalization to [0, 1] range is typically required for pre-trained models

Conclusion

Adding a batch dimension is essential when passing single images to TensorFlow models. Use np.newaxis or np.expand_dims() to add the required batch dimension, then pass the reshaped image to the model for prediction.