如何高效基于closest_colour返回值修改大型numpy RGBA数组元素？

我有一个RGBA格式的numpy数组（并非所有元素都是[0, 0, 0, 0]）：

array([[0, 0, 0, 0],
       [0, 0, 0, 0],
       [0, 0, 0, 0],
       ...,
       [0, 0, 0, 0],
       [0, 0, 0, 0],
       [0, 0, 0, 0]])

同时有一个closest_colour函数，能返回输入RGBA值最接近的5种颜色之一（绿色、红色、橙色、棕色、白色）：

def closest_colour(requested_colour):
    min_colours = {}
    for key, name in webcolors.CSS3_HEX_TO_NAMES.items():
        if name in ['green', 'red', 'orange', 'brown', 'white']:
            r_c, g_c, b_c = webcolors.hex_to_rgb(key)
            rd = (r_c - requested_colour[0]) ** 2
            gd = (g_c - requested_colour[1]) ** 2
            bd = (b_c - requested_colour[2]) ** 2
            min_colours[(rd + gd + bd)] = name
    return min_colours[min(min_colours.keys())]

尝试将函数应用到数组每个元素并修改对应元素时，使用了以下代码：

img_array[closest_colour(img_array) == 'green'] = (0, 255, 0, 1)
img_array[closest_colour(img_array) == 'red'] = (255, 0, 0, 1)
img_array[closest_colour(img_array) == 'brown'] = (92, 64, 51, 1)
img_array[closest_colour(img_array) == 'orange'] = (255, 165, 0, 1)
img_array[closest_colour(img_array) == 'white'] = (255, 255, 255, 0)

但出现错误：

TypeError: unhashable type: 'numpy.ndarray'

已知错误原因，但不知道针对大型图像数组的高效实现方式，求可行的高效方案。

原函数仅支持单个RGB值输入，直接传入numpy数组会触发类型错误，且逐元素循环处理大型数组效率极低。采用numpy向量化操作可以彻底解决这两个问题，所有计算在numpy底层C代码执行，速度提升显著。

步骤1：预定义目标颜色映射 #

先把匹配目标和替换值整理成numpy数组，避免重复查询webcolors：

import numpy as np
import webcolors

# 定义颜色匹配规则：名称 -> (参考RGB值, 替换用RGBA值)
color_map = {
    'green': (webcolors.hex_to_rgb('#008000'), (0, 255, 0, 1)),
    'red': (webcolors.hex_to_rgb('#ff0000'), (255, 0, 0, 1)),
    'brown': (webcolors.hex_to_rgb('#a52a2a'), (92, 64, 51, 1)),
    'orange': (webcolors.hex_to_rgb('#ffa500'), (255, 165, 0, 1)),
    'white': (webcolors.hex_to_rgb('#ffffff'), (255, 255, 255, 0))
}

# 提取为numpy数组，方便向量化计算
target_rgbs = np.array([rgb for rgb, _ in color_map.values()])
replace_rgbas = np.array([rgba for _, rgba in color_map.values()])

步骤2：向量化计算最接近颜色 #

利用numpy广播机制，一次性计算所有像素与目标颜色的距离，找到每个像素的最优匹配：

# 提取图像的RGB通道（原函数仅用RGB计算距离）
img_rgb = img_array[..., :3]

# 计算每个像素与所有目标颜色的距离平方，结果形状为(H, W, 5)
distances = np.sum((img_rgb[..., np.newaxis, :] - target_rgbs)**2, axis=-1)

# 找到每个像素距离最小的颜色索引，结果形状为(H, W)
closest_indices = np.argmin(distances, axis=-1)

步骤3：批量替换RGBA值 #

通过索引直接从替换数组中取值，完成整个数组的批量替换：

# 替换所有像素的RGBA值
img_array = replace_rgbas[closest_indices]

可选：保留全透明像素 #

如果需要保留原数组中[0,0,0,0]的全透明像素，可添加掩码判断：

# 生成全透明像素的掩码
transparent_mask = np.all(img_array == [0,0,0,0], axis=-1)

# 先执行替换，再恢复全透明像素
img_array = replace_rgbas[closest_indices]
img_array[transparent_mask] = [0,0,0,0]

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