Image Recognition using TensorFlow

Image Recognition using TensorFlow

Image recognition is a key application of deep learning, and TensorFlow, with its tf.keras API, provides powerful tools for building and training image recognition models. This chapter will guide you through creating a model to recognize images from a custom dataset organized in folders.

For this example, assume you have a dataset of images where each class of image is stored in its own sub-folder within a main dataset_image directory. For example: dataset_image/ class_A/ image1.jpg image2.png … class_B/ imageX.jpg imageY.jpeg … This is a common way to organize image datasets.

We will use tf.keras.utils.image_dataset_from_directory to load the images, which is a convenient utility for this type of folder structure. It automatically infers class labels from directory names and creates a tf.data.Dataset object, which is efficient for training.

# Step 1: Import necessary libraries
import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers
import matplotlib.pyplot as plt # For visualization (optional)
import numpy as np
import os # To simulate dataset creation if needed

# --- (Optional) Helper to create a dummy dataset for demonstration ---
def create_dummy_dataset(base_dir="dataset_image_demo", num_classes=2, img_per_class=10, img_size=(64,64)):
    if os.path.exists(base_dir):
        import shutil
        shutil.rmtree(base_dir) # Clean up if exists
    os.makedirs(base_dir, exist_ok=True)
    for i in range(num_classes):
        class_dir = os.path.join(base_dir, f"class_{i}")
        os.makedirs(class_dir, exist_ok=True)
        for j in range(img_per_class):
            # Create a dummy image (random noise)
            img_array = np.random.rand(img_size[0], img_size[1], 3) * 255
            img = keras.utils.array_to_img(img_array.astype('uint8'))
            img.save(os.path.join(class_dir, f"img_{j}.png"))
    print(f"Dummy dataset created at {base_dir}")
    return base_dir

# Create a dummy dataset if you don't have one
# Replace with your actual dataset path
IMAGE_DATASET_PATH = create_dummy_dataset(img_size=(64,64))
# IMAGE_DATASET_PATH = "path/to/your/dataset_image"

# --- Step 2: Load data using image_dataset_from_directory ---
IMG_WIDTH = 64
IMG_HEIGHT = 64
BATCH_SIZE = 32

# Create training dataset
train_dataset = keras.utils.image_dataset_from_directory(
    IMAGE_DATASET_PATH,
    labels='inferred',       # Infer labels from directory names
    label_mode='categorical', # For multi-class classification, yields one-hot encoded labels
    image_size=(IMG_HEIGHT, IMG_WIDTH),
    interpolation='nearest', # For resizing
    batch_size=BATCH_SIZE,
    shuffle=True,
    seed=42,
    validation_split=0.2,   # Use 20% of data for validation
    subset='training'       # Specify this is the training set
)

# Create validation dataset
validation_dataset = keras.utils.image_dataset_from_directory(
    IMAGE_DATASET_PATH,
    labels='inferred',
    label_mode='categorical',
    image_size=(IMG_HEIGHT, IMG_WIDTH),
    interpolation='nearest',
    batch_size=BATCH_SIZE,
    shuffle=False, # No need to shuffle validation data
    seed=42,
    validation_split=0.2,
    subset='validation'     # Specify this is the validation set
)

class_names = train_dataset.class_names
num_classes = len(class_names)
print(f"Found classes: {class_names}, Number of classes: {num_classes}")

# --- Step 3: Configure dataset for performance and normalize data ---
AUTOTUNE = tf.data.AUTOTUNE

train_dataset = train_dataset.cache().prefetch(buffer_size=AUTOTUNE)
validation_dataset = validation_dataset.cache().prefetch(buffer_size=AUTOTUNE)

# Normalize pixel values from [0, 255] to [0, 1]
normalization_layer = layers.Rescaling(1./255)

train_dataset = train_dataset.map(lambda x, y: (normalization_layer(x), y))
validation_dataset = validation_dataset.map(lambda x, y: (normalization_layer(x), y))

The data loading process reads images, resizes them, infers class labels, and splits the data into training and validation sets. The prefetch and cache methods optimize the data pipeline.

# --- Step 4: Define the CNN Model ---
model = keras.Sequential([
    layers.Input(shape=(IMG_HEIGHT, IMG_WIDTH, 3)), # Explicit input layer
    layers.Conv2D(32, (3, 3), padding='same', activation='relu', kernel_constraint=keras.constraints.MaxNorm(3)),
    layers.Dropout(0.2),
    layers.Conv2D(32, (3, 3), activation='relu', padding='same', kernel_constraint=keras.constraints.MaxNorm(3)),
    layers.MaxPooling2D(pool_size=(2, 2)),
    layers.Flatten(),
    layers.Dense(512, activation='relu', kernel_constraint=keras.constraints.MaxNorm(3)),
    layers.Dropout(0.5),
    layers.Dense(num_classes, activation='softmax') # num_classes inferred from data
])

# --- Step 5: Compile the model ---
# Using SGD optimizer as in the original example, but Adam is often a good default
epochs_train = 10 # Reduced for faster demo
lrate = 0.01
decay = lrate / epochs_train
sgd_optimizer = keras.optimizers.SGD(learning_rate=lrate, momentum=0.9, decay=decay, nesterov=False)

model.compile(loss='categorical_crossentropy',
              optimizer=sgd_optimizer,
              metrics=['accuracy'])

model.summary()

# --- Step 6: Train the model ---
print("\nTraining the model...")
# Add TensorBoard callback for monitoring (optional)
tensorboard_callback = keras.callbacks.TensorBoard(log_dir="./logs_image_rec", histogram_freq=1)

history = model.fit(
    train_dataset,
    validation_data=validation_dataset,
    epochs=epochs_train,
    callbacks=[tensorboard_callback] # Add callbacks here
)

# --- Step 7: Evaluate the model ---
print("\nEvaluating the model...")
loss, accuracy = model.evaluate(validation_dataset, verbose=0)
print(f"Validation Loss: {loss:.4f}")
print(f"Validation Accuracy: {accuracy*100:.2f}%")

# --- Step 8: Save the model (optional) ---
# Modern way to save the entire model (architecture, weights, optimizer state)
model.save("my_image_recognition_model.keras")
print("Saved model to disk as my_image_recognition_model.keras")

# Original way (saving architecture to JSON and weights to HDF5)
model_json = model.to_json()
with open("model_face_arch.json", "w") as json_file:
    json_file.write(model_json)
model.save_weights("model_face_weights.h5")
print("Saved model architecture to model_face_arch.json and weights to model_face_weights.h5")

# To load this type of saved model:
# from tensorflow.keras.models import model_from_json
# with open('model_face_arch.json', 'r') as json_file:
#     loaded_model_json = json_file.read()
# loaded_model = model_from_json(loaded_model_json)
# loaded_model.load_weights('model_face_weights.h5')
# print("Loaded model from disk")

The output will show the model summary, training progress, and final evaluation metrics. If you run this with the dummy dataset, accuracy will likely be around chance (e.g., 50% for 2 classes) as the images are random noise. With a real dataset, you’d expect meaningful accuracy.

This updated example uses modern TensorFlow 2.x features for data loading and model training, making the process more streamlined and efficient.

To run TensorBoard and visualize training, open a terminal and run: tensorboard --logdir ./logs_image_rec