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

Question

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

为了像下图一样可视化矩阵，我使用numpy.argmax将 3d 矩阵numpy.argmax成一个 2d 矩阵，该矩阵包含最可能类别的索引。 但我不知道如何继续获得我想要的可视化。

编辑

其实，我可以用一种微不足道的方式做到这一点。 也就是说，使用 for 循环遍历每个像素并为其分配颜色以获取图像。 但是，这不是矢量化编码，因为 numpy 具有加速矩阵操作的内置方法。 我需要为实时分段保存 CPU 周期。

Answer 1

这很容易。 您所需要的只是一个将 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转换为“颜色图”：

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