如何将RGB图像数组转换为plasma色标索引的整数数组？

核心思路 #

因为所有像素都来自plasma色标（共256种颜色），我们可以先建立RGB三元组→色标索引的映射，再用numpy的向量化操作批量转换整个图像，避免逐像素循环的低效问题。

具体步骤与代码 #

1. 将plasma色标转换为uint8格式的RGB数组
   matplotlib的cmap.colors返回0-1范围的浮点数，需先转换为和图像一致的uint8（0-255）格式：

   import matplotlib.pyplot as plt
   import numpy as np
   
   # 获取plasma色标并转换为uint8 RGB
   plasma_cmap = plt.get_cmap('plasma')
   plasma_rgb = (plasma_cmap.colors * 255).astype(np.uint8)
   

2. 建立RGB到索引的快速映射
   将每个RGB三元组转换为唯一整数键，再创建键到索引的字典：

   # 生成每个RGB三元组的唯一键（公式：R*256² + G*256 + B）
   plasma_keys = plasma_rgb[:, 0] * (256**2) + plasma_rgb[:, 1] * 256 + plasma_rgb[:, 2]
   # 创建键到索引的映射字典
   rgb_to_idx = {key: idx for idx, key in enumerate(plasma_keys)}
   

3. 批量转换图像像素为索引
   对整个图像执行相同的键转换，再通过字典映射得到索引数组：

   # 生成图像每个像素的唯一键
   img_keys = img[:, :, 0] * (256**2) + img[:, :, 1] * 256 + img[:, :, 2]
   # 转换为目标形状的索引数组
   idx_array = np.vectorize(rgb_to_idx.get)(img_keys)
   

更高效的numpy向量化实现（避免vectorize） #

如果图像尺寸较大，np.vectorize本质仍是循环，推荐用numpy广播匹配优化：

# 将plasma色标重塑为(256, 1, 1, 3)，方便和图像(1694,2012,3)广播匹配
plasma_rgb_reshaped = plasma_rgb[:, np.newaxis, np.newaxis, :]
# 计算每个像素与色标中所有颜色的匹配（布尔数组，形状(256,1694,2012)）
matches = np.all(img == plasma_rgb_reshaped, axis=-1)
# 获取每个像素对应的索引（每个像素仅匹配一个色标，argmax可直接得到对应位置）
idx_array = np.argmax(matches, axis=0)

验证示例 #

# 验证img[0,0]是否返回预期的0
print(idx_array[0,0])  # 输出0

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