voc 分割数据color彩图直接转灰度图

就是我有一批标注好的color图，那个类别用不同像素来表示。现在我只需要这批数据的4类。其他类别都不需要了。
任务就是把color图上所需要的类别转为灰度标签图。
其实也简单，知道每个类别对应的颜色了，那么就构造同等大小的灰度图，相应颜色就置为相应类别。举例如下：

比如猫😺的颜色是[255,0,128]，然后猫的类别为3。那么我就把彩图上面[255,0,128]像素在灰度图上对应位置置为3。
主要涉及python opencv 的位置像素操作。

import cv2
import os
import numpy as np

color_map = np.asarray([
    [180,120,120], #1  
    [224,5,255], #2
    [4,200,3], #3
    [6,230,230] #3
])


path_color_mark1 = "/data/color/"
path_save = "/data/save_gray/"

list_color = os.listdir(path_color_mark1)
for cnt,img_name in enumerate(list_color):
    print(cnt,img_name)
    path_img = path_color_mark1 + img_name
    img_color = cv2.imread(path_img)
    h,w,_ = img_color.shape

    len_color = len(color_map)
    gray_img = np.zeros([h, w, 1], np.uint8)
    for i in range(len_color):
        pos = img_color == color_map[i]
        pos_1 = pos[:, :, 0] ##由于pos的shape和color图一样是[h,w,c]。 pos里面的数值都是False,False,True  而gray图是当通道的，[h,w]

        label = i + 1
        if 2==i or 3==i:
            label = 3
        gray_img[pos_1] = label

    cv2.imwrite(path_save + img_name,gray_img)