# Image Classification with TensorFlow: Developing a Data Pipeline \(Part 1\)

- Canonical: https://33rdsquare.com/image-classification-with-tensorflow-developing-the-data-pipeline-part-1/
- Published: 2024-09-03
- Author: Jordan Brown
- Categories: [Artificial Intelligence & Machine Learning & ChatGPT](https://33rdsquare.com/category/tech/ai/)

---

Image classification is a fundamental task in computer vision that involves assigning a label or category to an input image. It has a wide range of applications, from identifying objects in self-driving cars to detecting diseases in medical scans. In recent years, deep learning techniques like convolutional neural networks (CNNs) have achieved state-of-the-art results on image classification benchmarks.

However, before you can train a CNN to classify images, you need to prepare your data. Real-world datasets are messy and complex. Images may have different sizes, encodings, and labels. This is where a data pipeline comes in. A data pipeline is a series of steps for transforming raw input data into a format suitable for training a machine learning model.

In this two-part guide, we‘ll walk through developing a data pipeline for image classification using TensorFlow, a popular open-source library for machine learning. Part 1 will focus on building the pipeline and training an initial CNN. In Part 2, we‘ll explore techniques for improving model performance and generalization.

Let‘s get started!

## The ETL Process

At a high level, a data pipeline can be broken down into three main stages:

1. Extract – retrieve the raw data from a source like a database or file system
2. Transform – preprocess and convert the data into a usable format
3. Load – feed the transformed data into a model for training or inference

This is known as the ETL process. While the details depend on your specific use case, ETL is a helpful framework for thinking about data pipelines.

One of the most tedious aspects of machine learning is writing the code for these ETL steps. Fortunately, many libraries now provide pre-built pipelines for common datasets and tasks. TensorFlow Datasets is one such library that we‘ll be using in this guide.

## TensorFlow Datasets

TensorFlow Datasets (TFDS) is a collection of ready-to-use datasets for TensorFlow. It handles downloading the data, preprocessing it into a standard format, and constructing a tf.data.Dataset object for input pipelines. TFDS has many built-in datasets across different domains, from image classification to natural language processing.

For this guide, we‘ll use the horses_or_humans dataset, which contains 300×300 RGB images of horses and humans. The goal is to train a binary classifier to distinguish between the two classes.

First, make sure you have TensorFlow and TensorFlow Datasets installed:

!pip install tensorflow tensorflow-datasets

Then import the required libraries:

```
import tensorflow as tf
import tensorflow_datasets as tfds
```

To load the horses_or_humans dataset:

```
train_ds, info = tfds.load(‘horses_or_humans‘, split=‘train‘, with_info=True, as_supervised=True)
val_ds, val_info = tfds.load(‘horses_or_humans‘, split=‘test‘, with_info=True, as_supervised=True)
```

The `load` function fetches the dataset and returns a `tf.data.Dataset` object. We specify the desired subset with the `split` argument (‘train‘ or ‘test‘). Setting `as_supervised=True` returns a tuple `(image, label)` for each example instead of a dictionary. The `with_info` argument provides access to metadata like the number of examples and classes.

```
print(info)
```

```
tfds.core.DatasetInfo(
    name=‘horses_or_humans‘,
    version=3.0.0,
    description=‘A dataset of 300x300 images of horses and humans.‘,
    homepage=‘https://laurencemoroney.com/datasets.html‘,
    features=FeaturesDict({
        ‘image‘: Image(shape=(300, 300, 3), dtype=tf.uint8),
        ‘label‘: ClassLabel(shape=(), dtype=tf.int64, num_classes=2),
    }),
    total_num_examples=1284,
    splits={
        ‘test‘: 256,
        ‘train‘: 1028,
    },
    supervised_keys=(‘image‘, ‘label‘),
    citation="""@ONLINE {horses_or_humans,
        author = "Laurence Moroney",
        title = "Horses or Humans Dataset",
        month = "feb",
        year = "2019",
        url = "http://laurencemoroney.com/datasets.html"
    }""",
    redistribution_info=,
)
```

We can see there are 1028 training examples and 256 validation examples. The images have shape (300, 300, 3) and labels are binary integers. The dataset is already split into train and test sets, so we don‘t need to manually partition it.

## Data Pipeline Transformations

Now that we‘ve loaded the dataset, let‘s apply some transformations. The two most common are shuffling and batching.

Shuffling randomizes the order of the examples, which helps prevent overfitting and promotes convergence during training. We can use the `shuffle` method to randomly shuffles the elements with a buffer size. A larger buffer size increases the randomness but takes more memory.

```
train_ds = train_ds.shuffle(1000)
```

Batching groups multiple examples into batches for efficiency. Instead of processing examples individually, the model can compute gradients on batches in parallel. Use the `batch` method to specify the batch size.

```
train_ds = train_ds.batch(32)
val_ds = val_ds.batch(32)
```

Here we‘ve shuffled the training set and batched both the training and validation sets with a batch size of 32. We can use method chaining to apply both transformations in a single line.

There are many other transformations you can add to your pipeline, like resizing images, normalizing pixel values, and data augmentation. We‘ll explore those in Part 2. For now, let‘s train an initial model.

## Training a CNN Classifier

With our data pipeline set up, we can build a CNN for image classification. We‘ll use Keras, a high-level API for building models in TensorFlow.

```
model = tf.keras.Sequential([
  tf.keras.layers.Conv2D(16, (3,3), activation=‘relu‘, input_shape=(300, 300, 3)),
  tf.keras.layers.MaxPooling2D(2, 2),
  tf.keras.layers.Conv2D(32, (3,3), activation=‘relu‘),
  tf.keras.layers.MaxPooling2D(2, 2),
  tf.keras.layers.Conv2D(64, (3,3), activation=‘relu‘),
  tf.keras.layers.MaxPooling2D(2, 2),
  tf.keras.layers.Conv2D(64, (3,3), activation=‘relu‘),
  tf.keras.layers.MaxPooling2D(2, 2),
  tf.keras.layers.Conv2D(64, (3,3), activation=‘relu‘),
  tf.keras.layers.MaxPooling2D(2, 2),
  tf.keras.layers.Flatten(),
  tf.keras.layers.Dense(512, activation=‘relu‘),
  tf.keras.layers.Dense(1, activation=‘sigmoid‘)
])

model.compile(optimizer=‘adam‘,
              loss=‘binary_crossentropy‘,
              metrics=[‘accuracy‘])

model.summary()
```

This defines a CNN architecture with five convolutional layers, each followed by max pooling, and two fully-connected layers at the end. The output layer has a single neuron with sigmoid activation for binary classification. We compile the model using the Adam optimizer and binary cross-entropy loss.

To train the model, simply pass the `Dataset` objects to the `fit` method:

```
epochs = 10
history = model.fit(train_ds, epochs=epochs, validation_data=val_ds)
```

After training for 10 epochs, let‘s evaluate performance by plotting the learning curves:

```
acc = history.history[‘accuracy‘]
val_acc = history.history[‘val_accuracy‘]

loss = history.history[‘loss‘]
val_loss = history.history[‘val_loss‘]

epochs_range = range(epochs)

plt.figure(figsize=(8, 8))
plt.subplot(1, 2, 1)
plt.plot(epochs_range, acc, label=‘Training Accuracy‘)
plt.plot(epochs_range, val_acc, label=‘Validation Accuracy‘)
plt.legend(loc=‘lower right‘)
plt.title(‘Training and Validation Accuracy‘)

plt.subplot(1, 2, 2)
plt.plot(epochs_range, loss, label=‘Training Loss‘)
plt.plot(epochs_range, val_loss, label=‘Validation Loss‘)
plt.legend(loc=‘upper right‘)
plt.title(‘Training and Validation Loss‘)
plt.show()
```

From the plots, we can see that training accuracy increases rapidly and reaches near 100%, while validation accuracy plateaus around 85%. The large gap between training and validation metrics is a telltale sign of overfitting. The model has learned patterns specific to the training set that don‘t generalize well to new data.

## Conclusion and Next Steps

In this guide, we walked through developing an image classification data pipeline with TensorFlow Datasets. We loaded the horses_or_humans dataset, applied transformations like shuffling and batching, and trained a CNN to classify images into two categories.

However, our initial model shows evidence of overfitting. In Part 2, we‘ll explore techniques for addressing this, including:

- Data augmentation: artificially increase the diversity of the training set by applying random transformations to images
- Regularization: add constraints to the model to reduce complexity and prevent memorization
- Transfer learning: leverage a model pretrained on a large dataset like ImageNet to improve performance on a smaller dataset

We‘ll also look at deploying our trained model and using it to make predictions on new images.

Image classification is a powerful technique with many real-world applications. I hope this guide has given you a practical starting point for building your own pipelines and models. Stay tuned for Part 2, and happy classifying!

---

Source: [Image Classification with TensorFlow: Developing a Data Pipeline \(Part 1\)](https://33rdsquare.com/image-classification-with-tensorflow-developing-the-data-pipeline-part-1/)