数据标注:不同类别的标注数据合并,并修改label文件?以下代码将为你一键搞定

# -*- coding: utf-8 -*-
# @Time : 2020/7/23 19:05
# @Author : for LongBo
# @File : tt.py
# @Software: PyCharm

import os

def get_classes_list(file_path):
    with open(file_path,'r') as f:
        lines = f.readlines()
    return lines

def labels_merge(source_dir1,source_dir2,new_classes):
    '''
    结合两个文件夹的label,更改两个文件夹中的label类别,并生成新的classes.txt文件
    :param source_dir1:待修改的第一个文件夹
    :param source_dir2:待修改的第二个文件夹
    :param new_classes:新的classes.txt全名
    :return:
    '''
    classes_file1 = os.path.join(source_dir1,'classes.txt')
    classes_file2 = os.path.join(source_dir2,'classes.txt')
    classes_list1 = get_classes_list(classes_file1)
    classes_list2 = get_classes_list(classes_file2)
    classes_merge = set(classes_list1 + classes_list2)
    classes_merge = list(classes_merge)

    print('...开始合并标签......')
    with open(new_classes,'w') as f:
        f.writelines(classes_merge)
    print('...合并标签完成......')

    print('...开始修改第一个文件夹......')
    for file in os.listdir(source_dir1):
        if not file.endswith('.txt') or file.startswith('classes'):
            continue
        current_txt = os.path.join(source_dir1,file)
        print(current_txt)
        with open(current_txt,'r') as f:
            lines = f.readlines()
        new_lines = []
        for line in lines:
            line_list = line.split(' ')
            line_list[0] = str(classes_merge.index(classes_list1[int(line_list[0])]))
            # print(line_list)
            new_lines.append(' '.join(line_list))
        print(new_lines)
        with open(current_txt,'w') as f:
            f.writelines(new_lines)
    print('...修改第一个文件夹完成......')

    print('...开始修改第二个文件夹......')
    for file in os.listdir(source_dir2):
        if not file.endswith('.txt') or file.startswith('classes'):
            continue
        current_txt = os.path.join(source_dir2, file)
        print(current_txt)
        with open(current_txt, 'r') as f:
            lines = f.readlines()
        new_lines = []
        for line in lines:
            line_list = line.split(' ')
            line_list[0] = str(classes_merge.index(classes_list2[int(line_list[0])]))
            # print(line_list)
            new_lines.append(' '.join(line_list))
        print(new_lines)
        with open(current_txt, 'w') as f:
            f.writelines(new_lines)
    print('...修改第二个文件夹完成......')

if __name__ == '__main__':
    labels_merge(r'D:\CV_Data\Region\SZ11Region',r'D:\CV_Data\Region\GZ5Region\merge',r'D:\CV_Data\Region\GZ5Region\classes.txt')

 

你可能感兴趣的:(图像处理,python,PyTorch)