Почему мои предсказания неверны при выполнении сегментации изображений с помощью TensorFlow?

Я пытаюсь выполнить семантическую сегментацию изображений с помощью TensorFlow. Чтобы просто заставить что-то работать, я беру одно учебное изображение, тренирую сеть на этом изображении некоторое время, а затем “тестирую” сеть на том же самом изображении, то есть я сильно переобучаюсь на этом одном изображении. Во время тестирования я ожидал бы, что сеть выдаст более или менее такие же сегментации, как и истинные данные, на которых она была обучена. Однако я получаю совершенно неверные прогнозы:

Я использую набор данных Synthia. Аннотации сегментации предоставляются в виде RGB изображений, где цвет каждого пикселя обозначает его класс:

Поскольку нет функции TensorFlow, которая автоматически берет RGB истинные данные и выводит их распределение классов (почему нет?), я просто использую значения цвета с плавающей точкой в качестве распределения классов истинных данных.

Почему мой код выдает такие произвольные карты прогнозирования, вместо того чтобы выдавать что-то, что более или менее напоминает соответствующее изображение аннотации истинных данных?

Мой код выглядит следующим образом:

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)

```

Вам следует использовать softmax в качестве активации выходного слоя, поскольку вы выполняете многоклассовую классификацию. Затем, для прогнозирования, вам нужно взять argmax выходного слоя

Также, так как вы используете категориальную кросс-энтропию, не забудьте преобразовать метки в one-hot кодирование

ИЗМЕНЕНИЕ ПОСЛЕ КОММЕНТАРИЕВ

Ваши метки предоставляются в виде RGB изображений, вы должны преобразовать их в подходящие метки, например, в скалярные целые значения.

Пример:

Оригинальные метки имеют размеры (H,W,3), а преобразованные метки имеют размеры (H,W,1). Если у вас 3 класса, у вас должно быть ровно 3 значения RGB, возьмем (255,0,0), (0,255,0) и (0,0,255). Таким образом, каждый пиксель с этими значениями должен быть преобразован в метку 0, 1 или 2 соответственно. Затем вы можете закодировать метки в виде матриц one-hot

Весь процесс выглядит следующим образом:

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