Tensorflow: Как создать изображение в стиле Pascal VOC

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

Например, предположим, что у меня есть сеть, которая тренируется с размером пакета 1 с 4 классами. Окончательные прогнозы сетей имеют форму [1, 3, 3, 4]

По сути, я хочу взять прогноз вывода и прогнать его через argmin, чтобы получить тензор, содержащий наиболее вероятный класс в каждой точке вывода:

[[[0, 1, 3], 
  [2, 0, 1],
  [3, 1, 2]]]

Аннотированные изображения используют цветовую палитру из 255 цветов для кодирования этикеток. У меня есть тензор, содержащий все тройки цветов:

  [[  0,   0,   0],
   [128,   0,   0],
   [  0, 128,   0],
   [128, 128,   0],
   [  0,   0, 128],
   ...
   [224, 224, 192]]

Как я могу получить тензор формы [1, 3, 3, 3] (одно цветное изображение 3x3), который индексируется в цветовой палитре с использованием значений, полученных из argmin?

[[palette[0], palette[1], palette[3]],
 [palette[2], palette[0], palette[1]],
 [palette[3], palette[1], palette[2]]]

Я мог бы легко обернуть некоторый numpy- и PIL-код в tf.py_func, но мне интересно, существует ли чистый способ Tensorflow для получения этого результата.

РЕДАКТИРОВАТЬ: Для тех, кому интересно, это решение, которое я получил, используя только numpy. Он работает неплохо, но мне все еще не нравится использование tf.py_func:

import numpy as np
import tensorflow as tf


def voc_colormap(N=256):
    bitget = lambda val, idx: ((val & (1 << idx)) != 0)

    cmap = np.zeros((N, 3), dtype=np.uint8)
    for i in range(N):
        r = g = b = 0
        c = i
        for j in range(8):
            r |= (bitget(c, 0) << 7 - j)
            g |= (bitget(c, 1) << 7 - j)
            b |= (bitget(c, 2) << 7 - j)
            c >>= 3

        cmap[i, :] = [r, g, b]
    return cmap


VOC_COLORMAP = voc_colormap()


def grayscale_to_voc(input, name="grayscale_to_voc"):
    return tf.py_func(grayscale_to_voc_impl, [input], tf.uint8, stateful=False, name=name)


def grayscale_to_voc_impl(input):
    return np.squeeze(VOC_COLORMAP[input])

person startup-sidious    schedule 22.03.2017    source источник


Ответы (1)


Вы можете использовать tf.gather_nd (), но вам нужно будет изменить формы палитры и логиты для получения желаемого изображения, например:

import tensorflow as tf
import numpy as np
import PIL.Image as Image

# We can load the palette from some random image in the PASCAL VOC dataset
palette = Image.open('.../VOC2012/SegmentationClass/2007_000032.png').getpalette()

# We build a random logits tensor of the requested size
batch_size = 1
height = width = 3
num_classes = 4
np.random.seed(1234)
logits = np.random.random_sample((batch_size, height, width, num_classes))
logits_argmax = np.argmax(logits, axis=3)  # shape = (1, 3, 3)
# array([[[3, 3, 0],
#         [1, 3, 1],
#         [0, 2, 0]]])

sess = tf.InteractiveSession()
image = tf.gather_nd(
    params=tf.reshape(palette, [-1, 3]),  # reshaped from list to RGB
    indices=tf.reshape(logits_argmax, [batch_size, -1, 1]))
image = tf.cast(tf.reshape(image, [batch_size, height, width, 3]), tf.uint8)
sess.run(image)
# array([[[[128, 128,   0],
#          [128, 128,   0],
#          [  0,   0,   0]],
#         [[128,   0,   0],
#          [128, 128,   0],
#          [128,   0,   0]],
#         [[  0,   0,   0],
#           [  0, 128,   0],
#           [  0,   0,   0]]]], dtype=uint8)

Полученный тензор можно напрямую передать в tf.summary.image (), но, в зависимости от вашей реализации, вы должны увеличить его перед сводкой.

person myagues    schedule 18.07.2017