Numpy数组求argmax后如何映射为(20,20,4)形状的RGBA颜色数组

你可以直接用Numpy的整数数组索引特性实现转换，核心代码非常简单：

rgba_arr = color_list[a]

你当前的a是形状为(20,20)的整数数组，每个元素取值范围为0~5，刚好对应color_list（形状为(6,4)，每一行是对应类别的RGBA值）的行索引。用a作为索引直接取color_list的值时，Numpy的广播机制会自动为每个位置的索引匹配对应的RGBA四元组，输出的rgba_arr形状自动就是(20,20,4)，完全符合需求。

完整可运行示例代码如下：

import numpy as np
import matplotlib.cm as cm

num_channels = 6
# 生成原始(20,20,6)数组
a = np.random.choice(range(100),(20,20,num_channels))
# 生成6个类别的RGBA颜色列表
color_list = cm.rainbow(np.linspace(0, 1, num_channels))
# 求最后一个轴的argmax，得到(20,20)的类别索引数组
a = a.argmax(axis=-1)
# 转换为RGBA数组
rgba_arr = color_list[a]
# 验证输出形状
print(rgba_arr.shape) # 输出 (20, 20, 4)

内容的提问来源于stack exchange，提问作者Carol Eisen