import os


# 类别名称列表
class_names = ['head', 'helmet', 'shoe']  # 替换为你的类别名称
target_class = ['head', 'helmet', 'shoe']  # 目标类别
target_class_id = {'head': 1, 'helmet':0, 'shoe':2}  # 目标类别的新 ID

def filter_and_rename_labels(label_directory, output_label_directory):
    if not os.path.exists(output_label_directory):
        os.makedirs(output_label_directory)  # 创建输出标签目录

    # 遍历标签目录中的所有文件
    for label_file in os.listdir(label_directory):
        if label_file.endswith('.txt'):  # 确保是标签文件
            with open(os.path.join(label_directory, label_file), 'r') as file:
                lines = file.readlines()

            # 过滤标签，只保留目标类别并修改类别 ID
            filtered_lines = []
            for line in lines:
                try:
                    class_id, x_center, y_center, box_width, box_height = map(float, line.split())
                    if class_names[int(class_id)] in target_class:
                        # 修改类别 ID 为 0
                        filtered_lines.append(f"{target_class_id[class_names[int(class_id)]]} {x_center} {y_center} {box_width} {box_height}\n")
                except:
                    print("标签文件错误", label_file)

            # 将过滤后的标签写入新的文件
            with open(os.path.join(output_label_directory, label_file), 'w') as file:
                file.writelines(filtered_lines)

# 示例用法
label_directory = '/home/admin-root/haotian/python哈汽锻8安全帽识别/export/duan8_11/labels_raw'  # 替换为你的标签文件夹路径
output_label_directory = '/home/admin-root/haotian/python哈汽锻8安全帽识别/export/duan8_11/labels'  # 替换为你希望保存过滤后标签的路径
filter_and_rename_labels(label_directory, output_label_directory)