根据颜色值将多维Numpy数组转换为二维数组

import numpy as np

# 定义颜色映射表
color_map = {
    (255, 0, 0): 0,  # 红色映射为索引0
    (0, 255, 0): 1,  # 绿色映射为索引1
    (0, 0, 255): 2   # 蓝色映射为索引2
}

# 多维Numpy数组
arr = np.array([
    [[255, 0, 0], [0, 255, 0]],
    [[0, 0, 255], [255, 0, 0]]
])

# 获取数组形状
shape = arr.shape

# 创建二维数组
new_arr = np.zeros((shape[0], shape[1]), dtype=int)

# 遍历多维数组
for i in range(shape[0]):
    for j in range(shape[1]):
        # 获取颜色值
        color = tuple(arr[i, j])
        # 查找颜色值对应的索引值
        index = color_map.get(color, -1)
        # 将索引值存储在二维数组中
        new_arr[i, j] = index

print(new_arr)