RGB彩色标签转单通道标签

代码： 

法一：

        在txt文件里读取数据。

from tqdm import tqdm
import numpy as np
import cv2
import os

# 标签中每个RGB颜色的值
VOC_COLORMAP = np.array([[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]])
# 标签其标注的类别
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']

# 处理txt中的对应图像
txt_path = '/home/DATA/database/VOC2012/VOCdevkit/VOC2012/ImageSets/Segmentation/trainval.txt' 
# 标签所在的文件夹
label_file_path = '/home/DATA/database/VOC2012/VOCdevkit/VOC2012/SegmentationClass' 
# 处理后的标签保存的地址
gray_save_path = '/home/DATA/database/Voc2012/new_label/'

with open(txt_path, 'r') as f:
    file_names = f.readlines()
    for name in tqdm(file_names):
        name = name.strip('\n')  # 去掉换行符
        label_name = name + '.png'  # label文件名
        label_url = os.path.join(label_file_path, label_name)

        mask = cv2.imread(label_url)
        mask = cv2.cvtColor(mask, cv2.COLOR_BGR2RGB) # 通道转换
        mask = mask.astype(int)  
        label_mask = np.zeros((mask.shape[0], mask.shape[1]), dtype=np.int16)
        # 标签处理
        for ii, label in enumerate(VOC_COLORMAP):
            locations = np.all(mask == label, axis=-1)
            label_mask[locations] = ii
        # 标签保存
        cv2.imwrite(gray_save_path+label_name, label_mask)

法二：

        直接在数据原始路径读取数据。

# -*- encoding: utf-8 -*-
'''
@description： 把图片数据从文件夹整理成csv文件，每一行代表其路径
'''
import numpy as np
import cv2
import os

# 标签中每个RGB颜色的值
VOC_COLORMAP = np.array([[0, 0, 0], [0, 0, 128]])

class rgb_gray(object):
    # 分割训练集 验证集 测试集
    # 做成对应的txt
    def __init__(self, label_dir, new_label_save_path):
        self.label_dir = label_dir
        self.new_label_save_path = new_label_save_path

    def read_path(self):
        for i, im in enumerate(os.listdir(self.label_dir)):
            label_name = im.split('.')[0] + '.png'  # 读取图片的名字，去label里面找，确保两个文件夹都有这个名字的图
            if os.path.exists(os.path.join(self.label_dir, label_name)):
                label_url = os.path.join(self.label_dir, label_name)
                mask = cv2.imread(label_url)
                mask = cv2.cvtColor(mask, cv2.COLOR_BGR2RGB)  # 通道转换
                mask = mask.astype(int)
                label_mask = np.zeros((mask.shape[0], mask.shape[1]), dtype=np.int16)
                # 标签处理
                for ii, label in enumerate(VOC_COLORMAP):
                    locations = np.all(mask == label, axis=-1)
                    label_mask[locations] = ii
                # 标签保存
                cv2.imwrite(self.new_label_save_path + "/" + label_name, label_mask)


if __name__ == "__main__":
    # ----注意：路径不能有中文----------
    # 标签所在上一层的路径
    label = r'D:\Desktop\02\labels'
    # 处理后的标签保存的地址
    new_label_save_path = r'D:\Desktop\02\01'
    tocsv = rgb_gray(label, new_label_save_path)
    tocsv.read_path()