Image Segmentation

What you’ll learn

the difference between semantic segmentation (label pixels by category) and instance segmentation (also separate individual object instances)
what a segmentation mask is, and how it’s just a label image
why segmentation models downsample with strided convolutions instead of pooling, then upsample with Conv2DTransposeConv2DTranspose
how to read a per-pixel prediction back into a mask with argmaxargmax

Beyond “what’s in this image”

So far, every CNN you’ve built answers one question per image: “is this a cat or a dog?” Image segmentation asks a much finer-grained question: “which pixels belong to the cat, and which belong to the background?” The output isn’t a single label anymore — it’s a full image, the same size as the input, where every pixel has been classified.

There are two flavors worth knowing:

Semantic segmentation — every pixel gets a category label like “cat” or “background.” If there are two cats in the picture, both sets of pixels get the same generic “cat” label.
Instance segmentation — goes further and tells the two cats apart: “cat 1” and “cat 2” get different labels, even though both are “cat” pixels.

This page focuses on semantic segmentation, using the classic foreground/background/contour setup from the Oxford-IIIT Pets dataset: every photo of a cat or dog comes with a segmentation mask — a same-sized image where each pixel is an integer: 11 (foreground), 22 (background), or 33 (contour).

diagram Segmentation model shape mermaid

  flowchart LR
  subgraph Encoder
    I["Input image
200x200x3"] --> C1["Conv2D stride 2
100x100x64"]
    C1 --> C2["Conv2D stride 2
50x50x128"]
    C2 --> C3["Conv2D stride 2
25x25x256"]
  end
  subgraph Decoder
    C3 --> T1["Conv2DTranspose stride 2
50x50x256"]
    T1 --> T2["Conv2DTranspose stride 2
100x100x128"]
    T2 --> T3["Conv2DTranspose stride 2
200x200x64"]
  end
  T3 --> O["Conv2D 1x1, softmax
200x200x3 (per-pixel classes)"]

An encoder compresses the image into small, information-dense feature maps; a decoder expands them back out to a full-resolution, per-pixel prediction.

Downsampling without losing “where”

The first half of a segmentation model looks a lot like a classification convnet: a stack of Conv2DConv2D layers with a growing number of filters. But there’s one crucial change. Classification convnets use MaxPooling2DMaxPooling2D to shrink feature maps, because they only care about whether a pattern is present anywhere in the image — max pooling’s translation invariance is a feature, not a bug, there.

Segmentation is the opposite: it needs to know exactly where every pixel’s class boundary is. A 2 × 2 max-pooling window keeps only the single largest value in each block and throws away which of the four positions it came from — that’s destructive to location information. So segmentation models downsample using strided convolutions instead: a Conv2DConv2D with strides=2strides=2 shrinks the feature map by the same factor a pooling layer would, but every output value is still a proper, spatially-anchored weighted sum of a small neighborhood, keeping location information intact.

segmentation_model.py

from tensorflow import keras
from tensorflow.keras import layers
 
def get_model(img_size, num_classes):
    inputs = keras.Input(shape=img_size + (3,))
    x = layers.Rescaling(1. / 255)(inputs)
 
    # --- Encoder: downsample with STRIDED convolutions, not pooling ---
    x = layers.Conv2D(64, 3, strides=2, activation="relu", padding="same")(x)
    x = layers.Conv2D(64, 3, activation="relu", padding="same")(x)
    x = layers.Conv2D(128, 3, strides=2, activation="relu", padding="same")(x)
    x = layers.Conv2D(128, 3, activation="relu", padding="same")(x)
    x = layers.Conv2D(256, 3, strides=2, activation="relu", padding="same")(x)
    x = layers.Conv2D(256, 3, activation="relu", padding="same")(x)
 
    # --- Decoder: upsample back to the original resolution ---
    x = layers.Conv2DTranspose(256, 3, activation="relu", padding="same")(x)
    x = layers.Conv2DTranspose(256, 3, activation="relu", padding="same", strides=2)(x)
    x = layers.Conv2DTranspose(128, 3, activation="relu", padding="same")(x)
    x = layers.Conv2DTranspose(128, 3, activation="relu", padding="same", strides=2)(x)
    x = layers.Conv2DTranspose(64, 3, activation="relu", padding="same")(x)
    x = layers.Conv2DTranspose(64, 3, activation="relu", padding="same", strides=2)(x)
 
    # One softmax score per pixel, per class
    outputs = layers.Conv2D(num_classes, 3, activation="softmax", padding="same")(x)
    return keras.Model(inputs, outputs)
 
model = get_model(img_size=(200, 200), num_classes=3)
model.compile(optimizer="rmsprop", loss="sparse_categorical_crossentropy")

segmentation_model.py

from tensorflow import keras
from tensorflow.keras import layers
 
def get_model(img_size, num_classes):
    inputs = keras.Input(shape=img_size + (3,))
    x = layers.Rescaling(1. / 255)(inputs)
 
    # --- Encoder: downsample with STRIDED convolutions, not pooling ---
    x = layers.Conv2D(64, 3, strides=2, activation="relu", padding="same")(x)
    x = layers.Conv2D(64, 3, activation="relu", padding="same")(x)
    x = layers.Conv2D(128, 3, strides=2, activation="relu", padding="same")(x)
    x = layers.Conv2D(128, 3, activation="relu", padding="same")(x)
    x = layers.Conv2D(256, 3, strides=2, activation="relu", padding="same")(x)
    x = layers.Conv2D(256, 3, activation="relu", padding="same")(x)
 
    # --- Decoder: upsample back to the original resolution ---
    x = layers.Conv2DTranspose(256, 3, activation="relu", padding="same")(x)
    x = layers.Conv2DTranspose(256, 3, activation="relu", padding="same", strides=2)(x)
    x = layers.Conv2DTranspose(128, 3, activation="relu", padding="same")(x)
    x = layers.Conv2DTranspose(128, 3, activation="relu", padding="same", strides=2)(x)
    x = layers.Conv2DTranspose(64, 3, activation="relu", padding="same")(x)
    x = layers.Conv2DTranspose(64, 3, activation="relu", padding="same", strides=2)(x)
 
    # One softmax score per pixel, per class
    outputs = layers.Conv2D(num_classes, 3, activation="softmax", padding="same")(x)
    return keras.Model(inputs, outputs)
 
model = get_model(img_size=(200, 200), num_classes=3)
model.compile(optimizer="rmsprop", loss="sparse_categorical_crossentropy")

Conv2DTranspose: convolution that upsamples

Conv2DTransposeConv2DTranspose is best thought of as the inverse of a strided convolution. If a Conv2D(128, 3, strides=2, padding="same")Conv2D(128, 3, strides=2, padding="same") shrinks a (100, 100, 64)(100, 100, 64) input down to (50, 50, 128)(50, 50, 128), then a matching Conv2DTranspose(64, 3, strides=2, padding="same")Conv2DTranspose(64, 3, strides=2, padding="same") takes that (50, 50, 128)(50, 50, 128) feature map and expands it back out to (100, 100, 64)(100, 100, 64) — the same spatial size you started with. Stack enough of these transposed convolutions and you can go from a small, information-dense (25, 25, 256)(25, 25, 256) bottleneck all the way back up to a full (200, 200, 3)(200, 200, 3) per-pixel prediction.

You can think of the first half of the model as compressing the image into a compact representation (each output location now summarizes a large chunk of the original picture), and the second half as decompressing that representation back into a full-resolution answer — one softmax probability distribution over classes, for every single pixel.

Loading the data and training the model

A segmentation mask is just an image, so with a small dataset like Oxford-IIIT Pets (7,390 photos, each with a matching trimap mask) it’s simplest to load everything straight into two NumPy arrays rather than build an input pipeline. Every input image and its mask are resized to the same img_sizeimg_size, and the raw mask labels (11, 22, 33) are shifted down by one so they line up with sparse_categorical_crossentropysparse_categorical_crossentropy’s expectation of zero-indexed classes:

load_segmentation_data.py

import os
import random
import numpy as np
from tensorflow.keras.utils import load_img, img_to_array
 
input_dir = "images/"
target_dir = "annotations/trimaps/"
img_size = (200, 200)
 
input_img_paths = sorted(
    os.path.join(input_dir, fname) for fname in os.listdir(input_dir) if fname.endswith(".jpg")
)
target_paths = sorted(
    os.path.join(target_dir, fname)
    for fname in os.listdir(target_dir)
    if fname.endswith(".png") and not fname.startswith(".")
)
 
# Shuffle both lists with the SAME seed so image/mask pairs stay lined up
random.Random(1337).shuffle(input_img_paths)
random.Random(1337).shuffle(target_paths)
 
def path_to_input_image(path):
    return img_to_array(load_img(path, target_size=img_size))
 
def path_to_target(path):
    img = img_to_array(load_img(path, target_size=img_size, color_mode="grayscale"))
    return img.astype("uint8") - 1   # shift {1, 2, 3} -> {0, 1, 2}
 
num_imgs = len(input_img_paths)
input_imgs = np.zeros((num_imgs,) + img_size + (3,), dtype="float32")
targets = np.zeros((num_imgs,) + img_size + (1,), dtype="uint8")
for i in range(num_imgs):
    input_imgs[i] = path_to_input_image(input_img_paths[i])
    targets[i] = path_to_target(target_paths[i])
 
# Hold out 1,000 samples for validation
num_val_samples = 1000
train_input_imgs, train_targets = input_imgs[:-num_val_samples], targets[:-num_val_samples]
val_input_imgs, val_targets = input_imgs[-num_val_samples:], targets[-num_val_samples:]

load_segmentation_data.py

import os
import random
import numpy as np
from tensorflow.keras.utils import load_img, img_to_array
 
input_dir = "images/"
target_dir = "annotations/trimaps/"
img_size = (200, 200)
 
input_img_paths = sorted(
    os.path.join(input_dir, fname) for fname in os.listdir(input_dir) if fname.endswith(".jpg")
)
target_paths = sorted(
    os.path.join(target_dir, fname)
    for fname in os.listdir(target_dir)
    if fname.endswith(".png") and not fname.startswith(".")
)
 
# Shuffle both lists with the SAME seed so image/mask pairs stay lined up
random.Random(1337).shuffle(input_img_paths)
random.Random(1337).shuffle(target_paths)
 
def path_to_input_image(path):
    return img_to_array(load_img(path, target_size=img_size))
 
def path_to_target(path):
    img = img_to_array(load_img(path, target_size=img_size, color_mode="grayscale"))
    return img.astype("uint8") - 1   # shift {1, 2, 3} -> {0, 1, 2}
 
num_imgs = len(input_img_paths)
input_imgs = np.zeros((num_imgs,) + img_size + (3,), dtype="float32")
targets = np.zeros((num_imgs,) + img_size + (1,), dtype="uint8")
for i in range(num_imgs):
    input_imgs[i] = path_to_input_image(input_img_paths[i])
    targets[i] = path_to_target(target_paths[i])
 
# Hold out 1,000 samples for validation
num_val_samples = 1000
train_input_imgs, train_targets = input_imgs[:-num_val_samples], targets[:-num_val_samples]
val_input_imgs, val_targets = input_imgs[-num_val_samples:], targets[-num_val_samples:]

With the data ready and get_model()get_model() from above compiled, training is a normal model.fit()model.fit() call — the only addition worth calling out is a ModelCheckpointModelCheckpoint callback that keeps only the best-scoring checkpoint by validation loss, since (like most models) this one overfits well before training finishes:

train_segmentation_model.py

import tensorflow as tf
 
model.compile(optimizer="rmsprop", loss="sparse_categorical_crossentropy")
 
callbacks = [
    tf.keras.callbacks.ModelCheckpoint("oxford_segmentation.keras", save_best_only=True),
]
history = model.fit(
    train_input_imgs, train_targets,
    epochs=50,
    callbacks=callbacks,
    batch_size=64,
    validation_data=(val_input_imgs, val_targets),
)

train_segmentation_model.py

import tensorflow as tf
 
model.compile(optimizer="rmsprop", loss="sparse_categorical_crossentropy")
 
callbacks = [
    tf.keras.callbacks.ModelCheckpoint("oxford_segmentation.keras", save_best_only=True),
]
history = model.fit(
    train_input_imgs, train_targets,
    epochs=50,
    callbacks=callbacks,
    batch_size=64,
    validation_data=(val_input_imgs, val_targets),
)

Plotting history.history["loss"]history.history["loss"] against history.history["val_loss"]history.history["val_loss"] typically shows validation loss bottoming out somewhere around epoch 25, then creeping back up as the model starts overfitting — exactly why the checkpoint is watching val_lossval_loss instead of just saving the final epoch’s weights. Reload the saved .keras.keras file and call model.predict()model.predict() on a held-out image to get the raw (200, 200, 3)(200, 200, 3) softmax output shown below.

Turning the prediction back into a mask

The model’s raw output has shape (200, 200, 3)(200, 200, 3) — three probabilities per pixel. To get a single mask, take the argmaxargmax along the last axis:

predict_mask.py

import numpy as np
 
# pred: shape (200, 200, 3) -- one softmax distribution per pixel
def display_mask(pred):
    mask = np.argmax(pred, axis=-1)   # (200, 200): the winning class per pixel
    mask *= 127                       # rescale {0, 1, 2} -> {0, 127, 254} for viewing
    return mask

predict_mask.py

import numpy as np
 
# pred: shape (200, 200, 3) -- one softmax distribution per pixel
def display_mask(pred):
    mask = np.argmax(pred, axis=-1)   # (200, 200): the winning class per pixel
    mask *= 127                       # rescale {0, 1, 2} -> {0, 127, 254} for viewing
    return mask

A model trained for around 50 epochs on the Oxford-IIIT Pets dataset overfits somewhere around epoch 25 (watch the validation loss to catch the best checkpoint), but produces solid foreground/background masks — with the odd artifact around sharp geometric shapes in the background.

Visualize it

Watch the model predict its mask one scanline at a time: the raw photo shows first, then a sweeping scan line fills in the foreground/background/contour prediction row by row, holds the finished mask for a moment, then fades back to the plain photo and starts over — exactly what argmaxargmax over the softmax output gives you, one row at a time:

sketch A predicted segmentation mask, overlaid on the image p5.js

A scan line sweeps down the image, filling in the predicted foreground/background/contour mask row by row, holds, then fades back to the plain photo. Click to restart.

Mini-checkpoint

Why do segmentation encoders downsample with strided convolutions instead of max pooling?

Max pooling keeps only the largest value in each window and forgets which position it came from — great for “is this pattern present anywhere?” classification, but destructive to the exact spatial information a per-pixel task depends on. Strided convolutions shrink the feature map the same way, without throwing that location information away.

🧪 Try It Yourself

Exercise 1 – Downsample With a Strided Convolution

Exercise 2 – Upsample Back With Conv2DTranspose

Exercise 3 – Turn Per-Pixel Scores Into a Mask

Continue to Interpreting What Convnets Learn (Grad-CAM) — open up the “black box” and see exactly which pixels convinced your convnet of its prediction.

If this helped you, consider buying me a coffee ☕

Buy me a coffee