如何使用Numpy将0/1矩阵值批量替换为对应RGB颜色数组

最优实现方案 #

直接用NumPy的整数数组索引即可一行完成映射，全程矢量化无循环，性能远高于手动for循环实现。

示例代码

import numpy as np

# 原有数据准备
class_colors = [[0,0,0], [0,255,0]]
# 注：np.random.random_integers 已弃用，推荐用randint生成随机整数
m = np.random.randint(0, 2, (5,5))

# 核心映射逻辑，仅需一行
color_matrix = np.array(class_colors)[m]

说明

1. 输出的color_matrix形状为(5, 5, 3)，恰好对应你需要的[宽度, 高度, RGB通道]格式，矩阵中原来的0/1值会自动替换为class_colors对应下标的颜色三元组。
2. 如果你更习惯用np.take，也可以替换为以下写法，效果完全一致：

color_matrix = np.take(np.array(class_colors), m, axis=0)

3. 该方法可以直接扩展到多类别场景：只要m中存储的是对应类别的下标，不管类别数是多少，都不需要修改映射逻辑。

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