python用opencv遍历RGB图片各个颜色的部分，生成mask

今天碰见了如下的图片，需要根据图片生成图片的mask：

下面分享一下我的代码：

import cv2
import numpy as np
import os, random, shutil
from collections import defaultdict

def generate_mask(fileName,pixel_rgb):
    img = cv2.imread(fileName)
    mask = np.zeros((512, 512), np.uint8)
    height = img.shape[0]
    weight = img.shape[1]
    channels = img.shape[2]
    unique_num=[]
    for row in range(height):            #遍历高
        for col in range(weight):         #遍历宽
            r=img[row, col, 0]
            g=img[row, col, 1]
            b=img[row, col, 2]
            
            if([r,g,b]==pixel_rgb):
                mask[row][col]=255
                # continue
            else:
                mask[row][col]=0
            if([r,g,b] !=[0,0,0]):
                t=str(r)+'_'+str(g)+'_'+str(b)
                unique_num.append(t)
    unique_num=list(set(unique_num))
    print(unique_num)
    pixel_rgb=[str(item) for item in pixel_rgb]
    cv2.imwrite('sketch_mask_{}.png'.format('_'.join(pixel_rgb)), mask)
# label_mapping('label_sketch.png')

color_list_p = [[0,0,204],[255,255,0],[204,0,0],[255,255,51],[0,255,255],[204,0,204],[0,153,76]]
for item in color_list_p:
    generate_mask('./WechatIMG360.png',item)

color_list_p就是在图片上取的颜色区域的rgb值。
部分生成的图片展示如下：