Tensorflow:如何创建Pascal VOC风格的图像

Question

我正在Tensorflow中实现一个语义分割网络，并试图弄清楚如何在训练期间写出标签的摘要图像。我希望以类似于Pascal VOC数据集中使用的class segmentation annotations的样式对图像进行编码。

例如，假设我有一个网络，它训练的批量大小为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]]

如何使用从argmin获得的值获得shape [1, 3, 3, 3] (单个3x3彩色图像)的张量，并将其索引到调色板中

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

我可以很容易地用tf.py_func包装一些numpy和PIL代码，但我想知道是否有一种纯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()，但您需要修改调色板和日志的形状以获得所需的图像，例如：

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