yolo格式数据集修改类别

import os


def modify_labels(folder_path):
    # 获取labels文件夹中所有的txt文件
    txt_files = [file for file in os.listdir(folder_path) if file.endswith('.txt')]

    for txt_file in txt_files:
        with open(os.path.join(folder_path, txt_file), 'r') as f:
            lines = f.readlines()

        # 将类别1替换为类别0
        new_lines = [modify_line(line) for line in lines]

        # 覆盖原来的txt文件内容
        with open(os.path.join(folder_path, txt_file), 'w') as f:
            f.writelines(new_lines)


def modify_line(line):
    class_id, *values = line.split()
    if class_id == '1': # 将类别1
        class_id = '0'   #修改为类别0
    new_line = class_id + " " + " ".join(values) + "\n"
    return new_line


if __name__ == "__main__":
    dataset_labels_folder = r"G:\datasets\beam\1000\txt_labels"
    modify_labels(dataset_labels_folder)