Why are my predictions broken when performing image segmentation with TensorFlow?

Question

I am attempting semantic image segmentation with TensorFlow. Just to get something working, I am taking this one training image, training the network on that image for a little while, and then "testing" the network on that same image, i.e. I am severely overfitting on that one image. During testing, I'd expect the network to yield more or less the same segmentations as the ground truth it was trained on. However, I get completely broken predictions:

I use the Synthia dataset. The segmentation annotation come in the form of RGB images, where each pixel's colour denotes its class:

Since there is no TensorFlow function which automatically takes RGB ground truth and deduces its class distribution from there (why not?), I simply use the float colour values as the ground truth class distribution.

Why does my code output such arbitrary prediction maps, instead of something that more or less resembles the corresponding ground truth annotation image?

My code is as follows:

from __future__ import absolute_import, division, print_function
import tensorflow as tf
import numpy as np
import os

tf.enable_eager_execution()

NUM_TRAINING_SAMPLES = 1
NUM_CLASSES = 3
BATCH_SIZE = 5
NUM_EPOCHS = 3
INPUT_SIZE = (256, 256, 3)

random_indices = np.random.choice(range(13000), NUM_TRAINING_SAMPLES)
directory_images = "C:/SYNTHIA/RGB"
directory_labels = "C:/SYNTHIA/GT"
train_images = np.array(os.listdir(directory_images))
train_labels = np.array(os.listdir(directory_labels))
train_images = train_images[random_indices]
train_labels = train_labels[random_indices]
train_images = [tf.read_file(os.path.join(directory_images, img)) for img in train_images]
train_labels = [tf.read_file(os.path.join(directory_labels, img)) for img in train_labels]
train_images = [tf.io.decode_image(img, channels=3) for img in train_images]
train_labels = [tf.io.decode_image(img, channels=3) for img in train_labels]
train_images = tf.image.resize_images(train_images, INPUT_SIZE[:2])
train_labels = tf.image.resize_images(train_labels, INPUT_SIZE[:2])
train_images = tf.image.convert_image_dtype(train_images, dtype=tf.uint8)
train_labels = train_labels / 256

train_dataset = tf.data.Dataset.from_tensor_slices((train_images, train_labels))
train_dataset = train_dataset.batch(1)
train_dataset = train_dataset.repeat()

def convolve(input_layer, num_filters):
    layer = tf.keras.layers.Conv2D(num_filters, (3, 3), padding='same')(input_layer)
    layer = tf.keras.layers.BatchNormalization()(layer)
    layer = tf.keras.layers.Activation('relu')(layer)
    layer = tf.keras.layers.Conv2D(num_filters, (3, 3), padding='same')(layer)
    layer = tf.keras.layers.BatchNormalization()(layer)
    layer = tf.keras.layers.Activation('relu')(layer)
    return layer

def downsample(input_layer, num_filters):
    layer = convolve(input_layer, num_filters)
    layer = tf.keras.layers.MaxPooling2D((2, 2), strides=(2, 2))(layer)
    return layer

def upsample(input_layer, num_filters):
    layer = tf.keras.layers.Conv2DTranspose(num_filters, (2, 2), strides=(2, 2), padding='same')(input_layer)
    layer = tf.keras.layers.BatchNormalization()(layer)
    layer = tf.keras.layers.Activation('relu')(layer)
    layer = tf.keras.layers.Conv2D(num_filters, (3, 3), padding='same')(layer)
    layer = tf.keras.layers.BatchNormalization()(layer)
    layer = tf.keras.layers.Activation('relu')(layer)
    layer = tf.keras.layers.Conv2D(num_filters, (3, 3), padding='same')(layer)
    layer = tf.keras.layers.BatchNormalization()(layer)
    layer = tf.keras.layers.Activation('relu')(layer)
    return layer

inputs = tf.keras.layers.Input(shape=INPUT_SIZE)  # 256
encoder0 = downsample(inputs, 32)  # 128
encoder1 = downsample(encoder0, 64)  # 64
encoder2 = downsample(encoder1, 128)  # 32
encoder3 = downsample(encoder2, 256)  # 16
encoder4 = downsample(encoder3, 512)  # 8
center = convolve(encoder4, 1024)  # center
decoder4 = upsample(center, 512)  # 16
decoder3 = upsample(decoder4, 256)  # 32
decoder2 = upsample(decoder3, 128)  # 64
decoder1 = upsample(decoder2, 64)  # 128
decoder0 = upsample(decoder1, 32)  # 256
outputs = tf.keras.layers.Conv2D(NUM_CLASSES, (1, 1), activation='sigmoid')(decoder0)

model = tf.keras.Model(inputs=[inputs], outputs=[outputs])  # model = tf.keras.Model(inputs=[images], outputs=[output])

model.compile(optimizer=tf.train.GradientDescentOptimizer(learning_rate=0.001), loss='categorical_crossentropy', metrics=['accuracy'])
#model.summary()
model.fit(train_dataset, batch_size=BATCH_SIZE, epochs=NUM_EPOCHS, steps_per_epoch=6, callbacks=[tf.keras.callbacks.TensorBoard(log_dir='./logs')])

#random_indices = np.random.choice(range(13000), NUM_TRAINING_SAMPLES)
test_images = np.array(os.listdir(directory_images))
test_images = test_images[random_indices]
test_images = [tf.read_file(os.path.join(directory_images, img)) for img in test_images]
test_images = [tf.io.decode_image(img, channels=3) for img in test_images]
test_images = tf.image.resize_images(test_images, INPUT_SIZE[:2])
test_images = tf.image.convert_image_dtype(test_images, dtype=tf.uint8)
test_dataset = tf.data.Dataset.from_tensors(test_images)

predictions = model.predict(test_dataset, batch_size=5, steps=1)
predictions = predictions * 256
predictions = tf.image.convert_image_dtype(predictions, dtype=tf.uint8)
jpg = tf.io.encode_jpeg(predictions[0])
tf.io.write_file("C:/SYNTHIA/prediction_map.jpg", jpg)

```

Answer 1

You should use softmax as the activation of the output layer, since you are performing multilclass classification. Then, for predicting, you have to take the argmax of the output layer

Aslo, since you are using categorical crossentropy, don't forget to convert the labels to one-hot encoding

EDIT AFTER THE COMMENTS

Your labels are provided as RGB images, you should convert it to appropriate labels, for instance, scalar integer values.

Example:

Original labels are of size (H,W,3) and the converted labels are of size (H,W,1). If you have 3 classes, you should have exactly 3 values of RGB, let's take (255,0,0), (0,255,0) and (0,0,255). Thus, each pixel having those values should be converted to label 0,1, or 2 accordingly. Then, you can encode the labels as one-hot matrixes

The whole pipe-line looks as follow:

• (255,0,0) --> 0 --> (1,0,0)
• (0,255,0) --> 1 --> (0,1,0)
• (0,0,255) --> 2 --> (0,0,1)