python-批量替换图像重指定的RGB值

import cv2
import os
import numpy as np
#color:            QR(1):128 0 0  MINIQR(2):0 128 0  DM(3):128 128 0  PDF(4):0 0 128
img_path='/data1/gyx/QR/multiyolov5_detect_seg/data/customdata/convert_tools/qr/segvoc/SegmentationClassPNG/'
out_img='/data1/gyx/QR/multiyolov5_detect_seg/data/customdata/convert_tools/qr/segvoc/mask/'

########### 第一种方法（循环遍历较慢） #########
# for im_name in os.listdir(img_path):
#     img=cv2.imread(img_path+im_name)
#     img = cv2.cvtColor(img,cv2.COLOR_BGR2RGB)
#     for i in range(img.shape[0]):
#         for j in range(img.shape[1]):
#             if img[i][j][0]==0 and img[i][j][1]==0 and img[i][j][2]==128:
#                 img[i][j]= 4,4,4
#             if img[i][j][0]==128 and img[i][j][1]==128 and img[i][j][2]==0:
#                 img[i][j]= 3,3,3
#             if img[i][j][0]==0 and img[i][j][1]==128 and img[i][j][2]==0:
#                 img[i][j]= 2,2,2
#             if img[i][j][0]==128 and img[i][j][1]==0 and img[i][j][2]==0:
#                 img[i][j]= 1,1,1
#     cv2.imwrite(out_img+im_name,img)

########### 第二种方法 （较快）##########
for im_name in os.listdir(img_path):
    im=cv2.imread(img_path+im_name)
    im = cv2.cvtColor(im,cv2.COLOR_BGR2RGB)
    im[np.all(im == (128, 0, 0), axis=-1)] = (1,1,1)
    im[np.all(im == (0, 128, 0), axis=-1)] = (2,2,2)
    im[np.all(im == (128, 128, 0), axis=-1)] = (3,3,3)
    im[np.all(im == (0, 0, 128), axis=-1)] = (4,4,4)
    cv2.imwrite(out_img+im_name,im)