代码:
法一:
在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()