Introduction

Pixel-wise classification for scene labeling is usually called semantic segmentation. Instead of predicting one label for the whole image, the model predicts one class for every pixel, producing a mask that has the same spatial structure as the input image.

That changes both the dataset format and the network design. A normal classifier collapses spatial detail, while a segmentation model must preserve or reconstruct it.

Start With Images And Masks

Each training example needs an image and a mask. The image is usually shape height x width x 3, while the mask stores one integer class ID per pixel with shape height x width.

In TensorFlow, it is important to preprocess masks differently from images. Images can be resized with interpolation, but segmentation masks should normally use nearest-neighbor resizing so class IDs do not get blurred into invalid intermediate values.

1import tensorflow as tf
2
3IMG_SIZE = (128, 128)
4
5
6def load_example(image_path, mask_path):
7    image = tf.io.read_file(image_path)
8    image = tf.image.decode_png(image, channels=3)
9    image = tf.image.resize(image, IMG_SIZE)
10    image = tf.cast(image, tf.float32) / 255.0
11
12    mask = tf.io.read_file(mask_path)
13    mask = tf.image.decode_png(mask, channels=1)
14    mask = tf.image.resize(mask, IMG_SIZE, method="nearest")
15    mask = tf.cast(mask, tf.int32)
16    mask = tf.squeeze(mask, axis=-1)
17
18    return image, mask

That gives you a clean pair of tensors for each example: a normalized image and an integer mask.

Use A Segmentation Architecture

A plain image classifier is the wrong model for pixel-wise scene labeling because it throws away most spatial information. You need a segmentation network such as U-Net, FCN, DeepLab, or some encoder-decoder design.

A small U-Net-style model is a good learning baseline:

1from tensorflow import keras
2from tensorflow.keras import layers
3
4
5def build_unet(num_classes: int):
6    inputs = keras.Input(shape=(128, 128, 3))
7
8    x1 = layers.Conv2D(32, 3, padding="same", activation="relu")(inputs)
9    x1 = layers.Conv2D(32, 3, padding="same", activation="relu")(x1)
10    p1 = layers.MaxPooling2D()(x1)
11
12    x2 = layers.Conv2D(64, 3, padding="same", activation="relu")(p1)
13    x2 = layers.Conv2D(64, 3, padding="same", activation="relu")(x2)
14    p2 = layers.MaxPooling2D()(x2)
15
16    b = layers.Conv2D(128, 3, padding="same", activation="relu")(p2)
17    b = layers.Conv2D(128, 3, padding="same", activation="relu")(b)
18
19    u1 = layers.UpSampling2D()(b)
20    u1 = layers.Concatenate()([u1, x2])
21    u1 = layers.Conv2D(64, 3, padding="same", activation="relu")(u1)
22
23    u2 = layers.UpSampling2D()(u1)
24    u2 = layers.Concatenate()([u2, x1])
25    u2 = layers.Conv2D(32, 3, padding="same", activation="relu")(u2)
26
27    outputs = layers.Conv2D(num_classes, 1, padding="same")(u2)
28    return keras.Model(inputs, outputs)

The final layer produces one logit vector per pixel, not one vector per image.

Compile With A Loss That Matches The Mask Format

If your masks store integer class IDs such as 0, 1, 2, and 3, a common loss is sparse categorical cross-entropy from logits.

1num_classes = 5
2model = build_unet(num_classes)
3
4model.compile(
5    optimizer="adam",
6    loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
7    metrics=["accuracy"]
8)

That works for single-class-per-pixel segmentation. If your task is binary segmentation, BinaryCrossentropy can also be appropriate depending on the mask encoding and output layer design.

Build The Dataset Pipeline Carefully

The image and mask must stay aligned through every transformation. That means random flips, crops, or rotations should be applied to both together.

1image_paths = tf.constant(["images/1.png", "images/2.png"])
2mask_paths = tf.constant(["masks/1.png", "masks/2.png"])
3
4dataset = tf.data.Dataset.from_tensor_slices((image_paths, mask_paths))
5dataset = dataset.map(load_example, num_parallel_calls=tf.data.AUTOTUNE)
6dataset = dataset.batch(8).prefetch(tf.data.AUTOTUNE)

For real projects you would split this into train and validation datasets and add augmentation before batching. The key is that mask geometry must always follow image geometry exactly.

Train And Decode Predictions

Training looks like ordinary Keras training:

history = model.fit(dataset, epochs=5)

To convert model output into predicted class IDs, take the argmax over the class axis:

1sample_images, _ = next(iter(dataset))
2logits = model.predict(sample_images)
3pred_mask = tf.argmax(logits, axis=-1)
4print(pred_mask.shape)

That gives one predicted class per pixel for each image in the batch.

Evaluate Beyond Accuracy

Pixel accuracy is easy to compute, but it can be misleading when one class dominates the image. A model can predict background very well and still fail badly on small but important classes such as pedestrians, signs, or lane markings.

For segmentation, metrics such as mean Intersection over Union and Dice score are often more informative. They tell you whether the predicted regions match the true regions, not just whether the model guessed the majority class often enough.

Common Pitfalls

One common mistake is resizing masks with bilinear interpolation, which corrupts class IDs. Another is using a plain classifier architecture and expecting it to produce a segmentation map. Developers also often mismatch loss and target format, for example by using sparse integer masks with a one-hot expectation or vice versa. Finally, augmentations frequently go wrong when they are applied to images but not to masks, which silently destroys training labels.

Summary

• Pixel-wise scene labeling in TensorFlow is a semantic-segmentation problem.
• Each example needs an image and an aligned per-pixel class mask.
• Use a segmentation architecture such as U-Net, not a plain classifier.
• Resize masks with nearest-neighbor logic and choose a loss that matches the mask format.
• Evaluate with segmentation-aware metrics because plain accuracy can hide poor class-level performance.