语义分割颜色标签（RGB）转方便训练的标签（灰度图）

 

#导入相关包
import matplotlib.pyplot as plt
import numpy as np
import glob
import scipy.misc
from PIL import Image
import re
#这是我的标签数据对应像素值
'''
label2color_dict = {
                    0:[0,0,0],
                    1:[128,0,0],
                    2:[0,128,150],
                    3:[128,128,0], 
                    4:[0,0,128],
                    5:[128,0,128],
                    6:[0,128,128],
                    7:[128,128,128],
                    8:[64,0,0],
                    9:[192,0,0],
                    10:[64,128,0],
                    11:[192,128,0],
                    12:[64,0,128],
                    13:[192,0,128],
                    14:[64,128,128],
                    15:[192,128,128],
                    16:[0,64,0],
                    17:[128,64,0],
                    18:[0,192,0],
                    19:[128,128,128]
                    }
'''
VOC_COLORMAP = [[0, 0, 0], [128, 0, 0], [0, 128, 0], [128, 128, 0],
                [0, 0, 128], [128, 0, 128], [0, 128, 128], [128, 128, 128],
                [64, 0, 0], [192, 0, 0], [64, 128, 0], [192, 128, 0],
                [64, 0, 128], [192, 0, 128], [64, 128, 128], [192, 128, 128],
                [0, 64, 0], [128, 64, 0], [0, 192, 0], [128, 192, 0],
                [0, 64, 128]]
zero = [0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19,20]
label2color_dict = dict(zip(zero, VOC_COLORMAP))
# VOC_COLORMAP相当于一个调色板
'''
VOC_CLASSES = [
    'background', 'aeroplane', 'bicycle', 'bird', 'boat', 'bottle', 'bus',
    'car', 'cat', 'chair', 'cow', 'diningtable', 'dog', 'horse', 'motorbike',
    'person', 'potted plant', 'sheep', 'sofa', 'train', 'tv/monitor']
'''


def get_keys(value):
#获取图片各像素value对应的key值

    p=0
    i=0
    
    for k,v in label2color_dict.items():
        if v == value:
            p = k
            break
        i += 1
    if i > 15:
        p = 0
    return p

#glob.glob方法获取全部文件名
img_path = glob.glob("/home/cjk/Downloads/VOCdevkit/VOC2012/SegmentationClass/*.png")
img_path = sorted(img_path, key=lambda img_path:int(re.search(r'(\d+)',img_path).group()))
for k in range(0, 10000):
    img = img_path[k]
    img = plt.imread(img)*255.0
    img = img[:,:,:3]
    img_label = np.zeros((img.shape[0],img.shape[1]),np.uint8)
    img_new_label = np.zeros((img.shape[0],img.shape[1]),np.uint8)
    for i in range(img.shape[0]):
        for j in range(img.shape[1]):
            value = list(img[i,j,:])
            img_label[i,j] = get_keys(value)
            #import pdb
            #pdb.set_trace()

            img_new_label[i,j] = img_label[i,j]

    label0 = Image.fromarray(np.uint8(img_new_label))
    p = str(k)
    scipy.misc.toimage(img_label).save("/home/cjk/Downloads/VOCdevkit/VOC2012/label_my/"+p+".jpg")