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

我正在使用神经网络进行语义分割（人类解析），比如拍摄一张人的照片作为输入，神经网络告诉我们每个像素最有可能是人类的头部、腿部、背景或其他部分。该算法运行平稳，输出

numpy.ndarray

。阵列的形状为

（1,23600400）

，其中600*400是输入图像的分辨率，23是类别数。3d矩阵看起来像一个23层堆叠的2d矩阵，其中每个层使用一个浮点矩阵来判断每个像素是否属于该类别

为了使矩阵可视化，如下图所示，我使用

numpy.argmax

将3d矩阵压缩为2d矩阵，该矩阵保存最可能类别的索引。但我不知道如何继续获得我想要的可视化效果

编辑

---

事实上，我可以用一种简单的方式来做。也就是说，使用for循环遍历每个像素并为其指定颜色以获得图像。然而，这不是矢量化编码，因为numpy内置了加速矩阵操作的方法。我需要节省CPU周期，以便进行实时分段。

这相当简单。您只需将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()

或者，如果您仅出于显示目的对colormap感兴趣，可以使用的

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()