Я работаю над реализацией сети семантической сегментации в 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])