Home » Machine Learning » How to Load Images in Machine Learning Models

How to Load Images in Machine Learning Models

Every type of data can be fed into machine learning models and images are no exception. As the saying goes, a picture is worth a thousand words. It’s appropriate since a single image holds a lot of information for a model to process.

We’re going to load and preprocess images using TensorFlow framework in this tutorial. In order to make this process as simple as possible, we’ll use high-level Keras preprocessing utilities and layers.

But first of all, we’ll need to import all the necessary python libraries.

import numpy as np
import pathlib
import os
import tensorflow as tf
import tensorflow_datasets as tfds

Prepare images for a machine learning model

We’ll demonstrate this process with an example flowers dataset, which contains images of 5 different kinds of flowers. Furthermore, they are separated by folders, where each flower kind has its own respective folder.

The following code will download this dataset onto your hard drive.

dataset_url = "https://storage.googleapis.com/download.tensorflow.org/example_images/flower_photos.tgz"
archive = tf.keras.utils.get_file(origin=dataset_url, extract=True)
data_dir = pathlib.Path(archive).with_suffix('')

Preprocess the dataset

Next, we’re going to use tf.keras.utils.image_dataset_from_directory utility function. In order to get the dataset ready for a model to be able to consume it, we need to preprocess it first.

There are a couple of arguments this function requires to work, so we need to define them.

batch_size = 32
img_height = 180
img_width = 180

We’re also going to split this dataset into training and validation set. Furthermore, we’ll use 80/20 split, meaning 80% of the dataset will be for training and 20% for validation.

It’s considered good practice to use these splits, so you can see whether your model overfits on the training data or not.

train_ds = tf.keras.utils.image_dataset_from_directory(
  data_dir,
  validation_split=0.2,
  subset="training",
  seed=123,
  image_size=(img_height, img_width),
  batch_size=batch_size)

val_ds = tf.keras.utils.image_dataset_from_directory(
  data_dir,
  validation_split=0.2,
  subset="validation",
  seed=123,
  image_size=(img_height, img_width),
  batch_size=batch_size)

Next step in preprocessing is the standardization of the images, where we bring numeric values of each color channel in pixels do into range between 0 and 1.

We can achieve this by using tf.keras.layers.Rescaling layer. Since the values for pixel channels range between 0 and 255, we also define this maximum in this layer.

normalization_layer = tf.keras.layers.Rescaling(1./255)

Next and final step in this preprocessing section is to configure the dataset for optimal performance.

AUTOTUNE = tf.data.AUTOTUNE

train_ds = train_ds.cache().prefetch(buffer_size=AUTOTUNE)
val_ds = val_ds.cache().prefetch(buffer_size=AUTOTUNE)

Conclusion

In conclusion, loading images in machine learning can be quite a complex process. Therefore, we can use TensorFlow, with its high-level functionality to make things easier.

I hope this article helped you clear up how loading images in TensorFlow works and perhaps even inspire you to research machine learning even more.