【数据集处理】获取json文件夹所有类名、更改某些类名、删除某些类名

1. 实现效果

2. 代码

• getclass：获取json文件夹中所有类名
• delclass：删除json文件夹中的某些类名
• changeclass：更改json文件夹中某些类名

import json
import os


def getclass(rootdir):
    """
    得到该目录下所有json文件中的label，即分割类别、类别数
        in_dir: 存放原始图像的目录
        out_dir: 重命名图像文件的保存目录
    """
    assert os.path.exists(rootdir), f'\"{rootdir}\" not exists. getclass over!'

    labels = []
    for file in os.listdir(rootdir):
        with open(os.path.join(rootdir, file), 'r') as load_f:
            load_dict = json.load(load_f)
            objects = load_dict['shapes']

            for i in range(0, len(objects)):
                label = objects[i]['label']
                if label not in labels:
                    labels.append(label)

    return labels, len(labels)


# 删除某些标签，目前可以暂定删掉：['Railway', 'Instrument_box', 'Sign', 'PJB', 'Signal_light']
# 再犹豫：['Cable', 'Pipe']
# 要统一，调整的：['Support'] 因为它有大有小
def delclass(json_dir, clslist, out_dir=None):
    """
    实现删除一些指定的类名标注
    :param json_dir: 原始的json文件夹
    :param clslist: 需要删除的json文件夹
    :param out_dir: 输出的存放json文件夹，若未指定则默认为就地更改，即保存在原始的json文件夹
    """
    # ------------------前期准备------------------- #
    assert os.path.exists(json_dir), f'\"{json_dir}\" not exixts. changeJsonLabelName over!'

    # 如果没有指定out_dir则就地改变
    if out_dir is None:
        out_dir = json_dir
    else:  # 若指定、但不存在则创建文件夹
        if not os.path.exists(out_dir):
            os.makedirs(out_dir)

    if not isinstance(clslist, list):  # 如果不是列表（就一个单个字符串），则添加进列表
        clslist = [clslist]
    # ---------------------------------------------- #

    json_files = os.listdir(json_dir)

    for json_file in json_files:
        jsonfile = os.path.join(json_dir, json_file)
        json_out = os.path.join(out_dir, json_file)
        # 读单个json文件
        with open(jsonfile, 'r', encoding='utf-8') as jf:
            load_dict = json.load(jf)
            objects = load_dict['shapes']

            for i in range(len(objects) - 1, -1, -1):  # 倒序
                label = objects[i]['label']
                if label in clslist:
                    del objects[i]

        # 将替换后的内容写入原文件
        with open(json_out, 'w') as new_jf:
            json.dump(load_dict, new_jf, indent=2)
    print('delete name over!')


def changeclass(json_dir, clslist, newclslist, out_dir=None):
    """
    实现更改对应类名
    :param json_dir: 原始json文件文件夹
    :param clslist: 待修改的类名列表（可只传入一个str类型）
    :param newclslist: 对应的新类名列表
    :param out_dir: 存放修改后的json文件，如果没有指定则就地更改
    """
    # ------------------前期准备------------------- #
    assert os.path.exists(json_dir), f'\"{json_dir}\" not exixts. changeJsonLabelName over!'

    if not isinstance(clslist, list):  # 如果不是列表（就一个单个字符串），则添加进列表
        clslist = [clslist]

    if not isinstance(newclslist, list):  # 如果不是列表（就一个单个字符串），则添加进列表
        newclslist = [newclslist]

    if len(clslist) != len(newclslist):
        print(f'旧类名{clslist}与新类名{newclslist}的长度不同，无法继续操作')
        return

    # 如果没有指定out_dir则就地改变
    if out_dir is None:
        out_dir = json_dir
    else:  # 若指定、但不存在则创建文件夹
        if not os.path.exists(out_dir):
            os.makedirs(out_dir)
    # ---------------------------------------------- #

    json_files = os.listdir(json_dir)
    for json_file in json_files:
        jsonfile = os.path.join(json_dir, json_file)
        json_out = os.path.join(out_dir, json_file)
        # 读单个json文件
        with open(jsonfile, 'r', encoding='utf-8') as jf:
            load_dict = json.load(jf)
            objects = load_dict['shapes']

            # ---------------改变类名的核心代码--------------- #
            for obj in objects:
                label = obj['label']
                if label in clslist:
                    obj['label'] = newclslist[clslist.index(label)]
            # --------------------------------------------- #
        # 将替换后的内容写入原文件
        with open(json_out, 'w') as new_jf:
            json.dump(load_dict, new_jf, indent=2)
    print('change name over!')


if __name__ == '__main__':
    print('# ------------获取json文件夹中所有类别-------------- #')
    rootdir = r'D:\SoftWareInstallMenu\JetBrains\PycharmProjects\tunnel_datadeal\crackop\delclass\labels(modify)'
    print(getclass(rootdir))

    print('\n# -------------------删除某些类别------------------ #')
    clslist = ['Support', 'Bolt_hole', 'Grouting_hole']
    out_dir = r'D:\SoftWareInstallMenu\JetBrains\PycharmProjects\tunnel_datadeal\crackop\delclass\labels(modify)_out'
    delclass(rootdir, clslist, out_dir)
    print(getclass(out_dir))

    print('\n# ------------------改变某些类别名------------------- #')
    clslist = ['Support', 'Bolt_hole', 'Grouting_hole']
    newlist = ['Support_1', 'Bolt_hole_1', 'Grouting_hole_1']
    out_dir = r'D:\SoftWareInstallMenu\JetBrains\PycharmProjects\tunnel_datadeal\crackop\delclass\labels(modify)_out_1'
    changeclass(rootdir, clslist, newlist, out_dir)
    print(getclass(out_dir))