查找图片中的某一点的索引（不用for循环）,制作语义分割多分类标签

a=np.array([[[0,0,128],
             [0,128,128]]])
red=np.array([0,0,128])
print(a[:,:,0])
print(a[:, :, 1])
print(a[:, :, 2])

print(a[:,:,0]==red[0])
print(a[:, :, 1] == red[1])
print(a[:, :, 2] == red[2])
print(a == red)
print(np.sum(a == red,axis=-1))
c=np.sum(a == red, axis=-1)
print(np.where(c==3))
y,x=np.where(c==3)
print(y)
print(x)
print(np.stack((x,y),axis=-1))

def find_pic_index(img,array_list):
    img_sum = np.sum(img == array_list, axis=-1)
    y, x =np.where(img_sum == 3)
    index_list = np.stack((x, y), axis=-1)
    return index_list
img_path='./SegmentationClassPNG/rec_1_0_20151031DandongRapideyeprjclip.png'
img=cv2.imread(img_path)
print(img.shape)
h, w, _ = img.shape
mask = np.zeros((h, w), dtype=np.int32)
red = np.array([0, 0, 128])
yellow = np.array([0, 128, 128])
green = np.array([0, 128, 0])
blue = np.array([128, 0, 0])
back_gro = np.array([0, 0, 0])

red_list = find_pic_index(img, red)
yellow_list = find_pic_index(img, yellow)
green_list = find_pic_index(img, green)
blue_list = find_pic_index(img, blue)
back_list = find_pic_index(img, back_gro)

for i in red_list:
    mask[i[1], i[0]] = 0
for i in yellow_list:
    mask[i[1], i[0]] = 255
for i in green_list:
    mask[i[1], i[0]] = 2
for i in blue_list:
    mask[i[1], i[0]] = 3
for i in back_list:
    mask[i[1], i[0]] = 128
mask = np.expand_dims(mask, axis=-1)
mask = np.concatenate((mask, mask, mask), axis=-1)
cv2.imwrite('./1.jpg', mask)