解决np.where()处理二维numpy数组时的RGB维度广播错误

你的报错核心原因是：原二维数组map的形状是(512,512)，而你传入np.where的替换值是长度为3的一维数组[255,255,255]和[0,0,0]，numpy无法将这两种形状进行广播，因此抛出维度不匹配的错误。以下是几种可行的解决方法：

方法1：扩展条件数组维度（最简洁） #

给判断条件map > 0.4增加一个末尾维度，让它从(512,512)变为(512,512,1)，这样就能和长度为3的RGB数组触发广播机制：

import numpy as np

# 假设map是(512,512)的二维numpy数组
condition = map > 0.4
# 扩展维度，axis=-1表示在最后添加一维
new_map = np.where(np.expand_dims(condition, axis=-1), [255,255,255], [0,0,0])
# 最终new_map形状为(512,512,3)，符合RGB格式

方法2：构造三维条件数组 #

将二维的条件数组复制三次，生成与目标RGB数组形状一致的(512,512,3)数组，再使用np.where：

condition = map > 0.4
# 将条件数组转为三维，最后一维重复3次
condition_3d = condition[:, :, np.newaxis].repeat(3, axis=-1)
new_map = np.where(condition_3d, 255, 0)

方法3：直接赋值（更直观） #

先创建一个全黑的三维数组，再通过布尔索引给符合条件的位置赋值白色：

# 创建形状为(512,512,3)的全0数组，指定dtype为uint8（适配RGB图像类型）
new_map = np.zeros((*map.shape, 3), dtype=np.uint8)
# 利用布尔索引批量赋值
new_map[map > 0.4] = [255,255,255]

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