如何将类别矩阵可视化为 RGB 图像？

Question

我正在使用神经网络进行语义分割（人体解析），比如将人的照片作为输入，神经网络告诉每个像素最有可能是头部、腿部、背景或其他一些部分人类的。该算法运行平稳并给出 numpy.ndarray 作为输出。数组的形状为(1,23,600,400)，其中600*400为输入图像的分辨率，23为类别数。 3d 矩阵看起来像一个 23 层堆叠的 2d 矩阵，其中每一层使用 float 矩阵来告诉每个像素属于该类别的可能性。

为了可视化如下图所示的矩阵，我使用 numpy.argmax 将 3d 矩阵压缩为包含最可能类别索引的 2d 矩阵。但是我不知道如何继续获得我想要的可视化效果。

编辑

其实我可以通过简单的方式做到。即用for循环遍历每一个像素点，并为其分配颜色，得到一张图像。然而，这不是矢量化编码，因为 numpy 有内置的方法来加速矩阵操作。我需要为实时分割节省 CPU 个周期。

Answer 1

这很容易。您所需要的只是 lookup table 将 23 个标签映射到独特的颜色。最简单的方法是使用一个 23×3 的 numpy 数组，每行存储相应标签的 RGB 值：

import numpy as np
import matplotlib.pyplot as plt
lut = np.random.rand(23, 3)   # using random mapping - but you can do better
lb = np.argmax(prediction, axis=1)  # converting probabilities to discrete labels
rgb = lut[lb[0, ...], :]  # this is all it takes to do the mapping.
plt.imshow(rgb)
plt.show()

或者，如果您只对用于显示目的的颜色图感兴趣，您可以使用 plt.imshow 的 cmap 参数，但这需要您将 lut 转换为 "colormap":

from matplotlib.colors import LinearSegmentedColormap
cmap = LinearSegmentedColormap.from_list('new_map', lut, N=23)
plt.imshow(lb[0, ...], cmap=cmap)
plt.show()

如何将类别矩阵可视化为 RGB 图像？

How to visualize a matrix of categories as an RGB image?

python

visualization

numpy

image-segmentation

deep-learning

编辑