TensorFlow
image processing
data augmentation
computer vision
machine learning

Tensorflow How to randomly crop input images and labels in the same way?

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

Introduction

When an image and its label must stay aligned, you cannot crop them independently with separate random calls. The correct solution is to generate one random crop location and apply that exact location to both tensors. This matters especially for segmentation masks, depth maps, and other dense labels where even a one-pixel mismatch corrupts the training data.

Why Separate Random Crops Break the Dataset

Suppose you do this:

python
cropped_image = tf.image.random_crop(image, [128, 128, 3])
cropped_label = tf.image.random_crop(label, [128, 128, 1])

Even if the crop sizes match, those two calls choose offsets independently. The image crop and label crop will usually come from different regions, which destroys supervision quality.

So the rule is simple: one random decision, two synchronized crops.

A Robust Shared-Offset Approach

The most general pattern is:

  1. sample crop offsets once
  2. apply those offsets to both image and label
python
1import tensorflow as tf
2
3def random_crop_pair(image, label, crop_height, crop_width, seed):
4    image_shape = tf.shape(image)
5    max_offset_height = image_shape[0] - crop_height + 1
6    max_offset_width = image_shape[1] - crop_width + 1
7
8    seeds = tf.random.experimental.stateless_split(seed, num=2)
9    offset_height = tf.random.stateless_uniform(
10        shape=[],
11        seed=seeds[0],
12        minval=0,
13        maxval=max_offset_height,
14        dtype=tf.int32,
15    )
16    offset_width = tf.random.stateless_uniform(
17        shape=[],
18        seed=seeds[1],
19        minval=0,
20        maxval=max_offset_width,
21        dtype=tf.int32,
22    )
23
24    cropped_image = tf.image.crop_to_bounding_box(
25        image, offset_height, offset_width, crop_height, crop_width
26    )
27    cropped_label = tf.image.crop_to_bounding_box(
28        label, offset_height, offset_width, crop_height, crop_width
29    )
30
31    return cropped_image, cropped_label
32
33
34image = tf.random.uniform([256, 256, 3])
35label = tf.ones([256, 256, 1], dtype=tf.int32)
36
37cropped_image, cropped_label = random_crop_pair(
38    image, label, 128, 128, seed=tf.constant([123, 456], dtype=tf.int32)
39)
40
41print(cropped_image.shape)
42print(cropped_label.shape)

This pattern is reliable because the offsets are shared explicitly.

Why Stateless Randomness Helps

TensorFlow has both stateful and stateless random APIs. For paired augmentations, stateless randomness is especially useful because the same seed always produces the same random choice. That makes debugging and reproducibility much easier.

If you are building a tf.data pipeline, passing a seed or deriving one per example can keep augmentation deterministic when needed.

Concatenation as a Shortcut

If image and label have compatible spatial shapes and you are comfortable combining them temporarily, another trick is concatenating them along the channel axis and cropping once.

python
1import tensorflow as tf
2
3image = tf.random.uniform([256, 256, 3], dtype=tf.float32)
4label = tf.cast(tf.ones([256, 256, 1]), tf.float32)
5
6combined = tf.concat([image, label], axis=-1)
7cropped = tf.image.stateless_random_crop(
8    combined,
9    size=[128, 128, 4],
10    seed=[7, 11],
11)
12
13cropped_image = cropped[:, :, :3]
14cropped_label = tf.cast(cropped[:, :, 3:], tf.int32)

This is compact, but it is slightly less general because image and label dtypes often differ in real datasets. The shared-offset approach is more robust when labels are integer masks.

Using the Pattern in a tf.data Pipeline

Here is how the shared crop can fit into dataset preprocessing:

python
1def preprocess(image, label):
2    image = tf.image.resize(image, [256, 256])
3    label = tf.image.resize(label, [256, 256], method="nearest")
4
5    seed = tf.constant([1, 2], dtype=tf.int32)
6    image, label = random_crop_pair(image, label, 128, 128, seed)
7
8    return image, label

Notice the label resize uses nearest-neighbor interpolation. For masks and class labels, you usually do not want bilinear interpolation because it invents intermediate values.

Segmentation Masks Need Special Care

For image classification, a crop mismatch might just reduce accuracy. For segmentation, it directly corrupts the label. That is why synchronized transforms are non-negotiable for:

  • segmentation masks
  • depth maps
  • optical flow
  • keypoint heatmaps
  • paired image-to-image translation targets

The same principle applies to flips, rotations, and resizing.

Common Pitfalls

One common mistake is calling tf.image.random_crop separately on the image and label and assuming equal size means equal crop location. It does not.

Another issue is resizing masks with bilinear interpolation before cropping. That can blur or invent class values. Use nearest-neighbor for label-like tensors.

Developers also sometimes use random augmentation layers independently on image and label. Unless the randomness is shared, the pair will drift apart.

Finally, if reproducibility matters, prefer stateless random APIs or explicit seeds rather than relying on global randomness and hoping the two calls stay synchronized.

Summary

  • Image and label crops must share the exact same random offsets.
  • The safest pattern is to sample offsets once and crop both tensors with crop_to_bounding_box.
  • Stateless random APIs are useful for reproducible paired augmentation.
  • Concatenation and one crop call can work when dtypes and shapes make that convenient.
  • For masks and labels, keep interpolation and augmentation choices label-safe.

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