Practice Now
TensorFlow
image processing
machine learning
data labeling
image classification

Tensorflow read images with labels

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design
TensorFlow & KerasComputer VisionMachine Learning

Understanding TensorFlow's Data Pipeline for Images with Labels

TensorFlow, an open-source library developed by Google, is a powerful tool for machine learning and neural network projects. A crucial part of any machine learning workflow involves efficiently loading and processing data. Images paired with corresponding labels are commonly used inputs for training models in TensorFlow. This article explains how to read and handle such data using TensorFlow, along with examples and technical explanations.

1. Preparing the Dataset

When working with images and labels in TensorFlow, it's common to have images stored in directories where each directory corresponds to a specific label. For instance, in an image classification task, you might have a folder named "cats" containing cat images and another folder "dogs" containing dog images.

2. TensorFlow Datasets (TFDS)

TensorFlow Datasets (TFDS) is a collection of ready-to-use datasets for various machine learning tasks. It offers a convenient way to load datasets for training and evaluation. However, when dealing with custom data, one often uses TensorFlow's low-level data handling APIs.

3. Loading Images with Labels

To load images and their corresponding labels in TensorFlow, you can use the tf.data API, which is both efficient and scalable. Below are the steps to read images along with labels:

Step 1: Setting Up Directory Structure

Assume your data is structured in a hierarchy like the following:

 
1data/
2  train/
3    cats/
4      cat001.jpg
5      cat002.jpg
6    dogs/
7      dog001.jpg
8      dog002.jpg
9  validation/
10    cats/
11    dogs/

Step 2: Import Libraries

python
import tensorflow as tf
import os

Step 3: Data Preprocessing Function

python
1def process_image(file_path):
2    # Load the raw data from the file as a string
3    img = tf.io.read_file(file_path)
4    # Decode it into a dense tensor
5    img = tf.image.decode_jpeg(img, channels=3)
6    # Normalize image data to 0–1 range
7    img = tf.image.convert_image_dtype(img, tf.float32)
8    # Resize to a desired size
9    img = tf.image.resize(img, [128, 128])
10    return img

Step 4: Create Dataset

python
1def load_data(data_dir):
2    # List all the directories and make class_names based on the folder names
3    class_names = sorted([d.name for d in os.scandir(data_dir) 
4                          if d.is_dir()])
5    class_indices = {name: index for index, name in enumerate(class_names)}
6    
7    # Gather file paths and corresponding labels
8    file_paths = []
9    labels = []
10    for label, class_name in enumerate(class_names):
11        dir_path = os.path.join(data_dir, class_name)
12        for filename in os.listdir(dir_path):
13            file_paths.append(os.path.join(dir_path, filename))
14            labels.append(label)
15    
16    # Convert to Tensor instances
17    file_paths = tf.constant(file_paths)
18    labels = tf.constant(labels)
19
20    # Create Dataset from tensors
21    dataset = tf.data.Dataset.from_tensor_slices((file_paths, labels))
22    return dataset

Step 5: Map and Batch the Dataset

python
1def prepare_data(data_dir):
2    dataset = load_data(data_dir)
3    
4    # Use map to preprocess and load data
5    dataset = dataset.map(lambda x, y: (process_image(x), y), 
6                          num_parallel_calls=tf.data.AUTOTUNE)
7    # Shuffle, repeat, and batch dataset
8    dataset = dataset.shuffle(buffer_size=1000)
9    dataset = dataset.batch(32)
10    dataset = dataset.cache()
11    dataset = dataset.prefetch(buffer_size=tf.data.AUTOTUNE)
12    
13    return dataset

4. Example of Using the Dataset

python
1train_dataset = prepare_data('data/train')
2validation_dataset = prepare_data('data/validation')
3
4# Iterate through the batches of train_dataset
5for images, labels in train_dataset.take(1):
6    print(images.shape)  # Batch of images
7    print(labels.shape)  # Corresponding labels

5. Summary Table of Key Points

StepDescriptionExample Code Snippet
Directory StructureOrganize images into directories named by their labelsdata/train/cats/cat001.jpg
Import LibrariesEnsure TensorFlow and other necessary libraries are importedimport tensorflow as tf
Data PreprocessingFunction to read and preprocess imagestf.image.resize(img, [128, 128])
Create DatasetLoad data and convert into a TensorFlow datasettf.data.Dataset.from_tensor_slices
Map & BatchPreprocess, shuffle, and batch the datasetdataset.batch(32)

Additional Tips

  • Optimizing Data Input: Use tf.data.AUTOTUNE to automatically adjust to optimally utilize resources.
  • Augmentation: For better model generalization, consider including image augmentation techniques using tf.image module functions.
  • Monitoring Dataset Creation: Check the dataset size and shapes to confirm the correctness of data ingestion.

This guide highlights the main steps involved in loading and preprocessing images with labels in TensorFlow, enabling efficient data input pipelining crucial for training machine learning models.

Related reading
Free course
Beginner
7 lessons
2 hours
Tackling System Design Interview Problems

A short course that equips you with the skills to approach system design interviews methodically.

Start the free course
Track what you have practised

A free account saves your progress, solutions and study plan across every problem on Codemia.

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design
Codemia
© 2026 Codemia
Resources
Legal
Social
All Rights Reserved.