如何在TensorFlow中将多类one-hot张量转换为RGB？

Question

我有一个形状为[None, 128, 128, n_classes] 的张量。这是一个one-hot tensor，其中最后一个索引包含多个类的分类值（总共有n_classes）。 在实践中，最后一个通道具有指示每个像素类别的二进制值：例如当一个像素在通道 C 中有 1 时，表示它属于 C 类；此像素在其他地方将有 0。

现在，我希望将这个单热张量转换为 RGB 图像，我想在 Tensorboard 上进行绘制。每个类都必须与不同的颜色相关联，以便于解释。

你知道怎么做吗？

谢谢，G。

---

编辑 2：

答案中添加了解决方案。

---

编辑 1：

我当前的实现（不工作）：

def from_one_hot_to_rgb(incoming, palette=None):
    """ Assign a different color to each class in the input tensor """
    if palette is None:
        palette = {
            0: (0, 0, 0),
            1: (31, 12, 33),
            2: (13, 26, 33),
            3: (21, 76, 22),
            4: (22, 54, 66)
        }

    def _colorize(value):
        return palette[value]

    # from one-hot to grayscale:
    cmap = tf.expand_dims(tf.argmax(incoming, axis=-1), axis=-1)

    # flatten input tensor (pixels on the first axis):
    B, W, H, C = get_shape(camp)  # this returns batch_size, 128, 128, 5
    cmap_flat = tf.reshape(cmap, shape=[B * W * H, C])

    # assign a different color to each class:
    cmap = tf.map_fn(lambda pixel:
                     tf.py_func(_colorize, inp=[pixel], Tout=tf.int64),
                     cmap_flat)

    # back to original shape, but RGB output:
    cmap = tf.reshape(cmap, shape=[B, W, H, 3])

    return tf.cast(cmap, dtype=tf.float32)

Answer 1

我会使用 matplotlib 中的 imshow* 或 matshow* 创建绘图，然后使用 this answer 或同一问题的其他答案将其显示在张量板上。

import matplotlib.pyplot as plt

plt.imshow(tf.argmax(imgs[0], axis=-1))

这种方法的一个优点是您不必担心类到颜色的映射。

---

要修复您已有的代码，首先您应该注意传递给 colorize 的参数是长度为 1 的 numpy 数组，而不是 int；这不是可散列的，因此不能用于字典键。您可以将其转换为 int 类型，就像 palette[int(value)] 一样。

我在这里和那里更改了您代码中的一些内容，并在大小为 1 的随机批次上对其进行了测试，最终代码如下所示：

def from_one_hot_to_rgb(incoming, palette=None):
    """ Assign a different color to each class in the input tensor """
    if palette is None:
        palette = {i: tf.constant(color, dtype='int64') for i, color in enumerate(
            ((0, 0, 0),
            (31, 12, 33),
            (13, 26, 33),
            (21, 76, 22),
            (22, 54, 66))
        )}

    # from one-hot to grayscale:
    B, W, H, _ = incoming.get_shape()   # this returns batch_size, 128, 128, 5
    cmap = tf.reshape(tf.argmax(incoming, axis=-1), [-1, 1])
    cmap = tf.map_fn(lambda value: palette[int(value)], cmap)

    # back to original shape, but RGB output:
    cmap = tf.reshape(cmap, shape=[B, W, H, 3])

    return tf.cast(cmap, dtype=tf.float32)

Answer 2

解决方案 1（慢）

一个可能的解决方案，类似于初始代码如下。请注意，由于 TensorFlow tf.map_fn 的已知 problem，这可能会非常慢

def from_one_hot_to_rgb_bkup(incoming, palette=None):

    if palette is None:
        palette = {i: tf.constant(color, dtype='int64') for i, color in enumerate(
            ((0, 0, 0),
            (31, 12, 33),
            (13, 26, 33),
            (21, 76, 22),
            (22, 54, 66))
        )}

    # from one-hot to grayscale:
    B, W, H, _ = get_shape(incoming)
    gray = tf.reshape(tf.argmax(incoming, axis=-1, output_type=tf.int32), [-1, 1], name='flatten')

    # assign colors to each class
    rgb = tf.map_fn(lambda pixel:
                    tf.py_func(lambda value: palette[int(value)], inp=[pixel], Tout=tf.int32),
                    gray, name='colorize')

    # back to original shape, but RGB output:
    rgb = tf.reshape(rgb, shape=[B, W, H, 3], name='back_to_rgb')

    return tf.cast(rgb, dtype=tf.float32)

解决方案 2（快速）

基于this 的回答，更快的解决方案可以使用tf.gather：

def from_one_hot_to_rgb_bkup(incoming, palette=None):

    if palette is None:
        palette = {i: tf.constant(color, dtype='int64') for i, color in enumerate(
            ((0, 0, 0),
            (31, 12, 33),
            (13, 26, 33),
            (21, 76, 22),
            (22, 54, 66))
        )}

    _, W, H, _ = get_shape(incoming)
    palette = tf.constant(palette, dtype=tf.uint8)
    class_indexes = tf.argmax(incoming, axis=-1)

    class_indexes = tf.reshape(class_indexes, [-1])
    color_image = tf.gather(palette, class_indexes)
    color_image = tf.reshape(color_image, [-1, W, H, 3])

    color_image = tf.cast(color_image, dtype=tf.float32)