Tensorflow:如何创建Pascal VOC样式图像

Question

我正在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])

Answer 1

您可以使用tf.gather_nd()，但您需要修改调色板和 logits 的形状以获得所需的图像，例如：

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()，但根据您的实现，您应该在摘要之前对其进行上采样。