删除所有多余的yolo格式下txt文件标签，只保留个别类别

读取每个YOLO标注文件，检查每个标签是否属于你希望保留的类别列表，如果不在列表中，则不会写入到新文件中。

def filter_yolo_annotations(input_dir, output_dir, classes_to_keep):
    """
    过滤YOLO格式的标注文件，只保留指定的类别。

    :param input_dir: 包含原始标注文件的目录路径。
    :param output_dir: 保存过滤后标注文件的目录路径。
    :param classes_to_keep: 需要保留的类别索引列表。
    """
    # 确保输出目录存在
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)

    # 遍历输入目录中的所有文件
    for filename in os.listdir(input_dir):
        if filename.endswith('.txt'):
            input_path = os.path.join(input_dir, filename)
            output_path = os.path.join(output_dir, filename)
            
            # 初始化新文件的内容列表
            new_file_content = []
            
            with open(input_path, 'r') as file:
                for line in file:
                    parts = line.strip().split()
                    class_index = int(parts[0])
                    # 检查类别索引是否在保留列表中
                    if class_index in classes_to_keep:
                        # 如果是，添加到新文件内容列表
                        new_file_content.append(line)
            
            # 将过滤后的内容写入新文件
            with open(output_path, 'w') as file:
                file.writelines(new_file_content)

            print(f"Filtered file saved to {output_path}")

# 指定包含YOLO标注文件的目录路径
input_dir = 'path/to/yolo/annotations'
# 指定保存过滤后标注文件的目录路径
output_dir = 'path/to/filtered/annotations'
# 指定需要保留的类别索引列表
classes_to_keep = [1, 3, 5]  # 例如，只保留类别索引为1, 3, 5的标注

# 调用函数
filter_yolo_annotations(input_dir, output_dir, classes_to_keep)

• 函数定义：filter_yolo_annotations函数接受三个参数：原始标注文件所在的目录路径、输出过滤后标注文件的目录路径和需要保留的类别索引列表。
• 创建输出目录：如果输出目录不存在，使用os.makedirs创建它。
• 遍历文件：遍历输入目录中的所有文件，找到扩展名为.txt的标注文件。
• 过滤标注：对于每个文件，读取每一行，检查类别索引是否在classes_to_keep列表中。
• 写入新文件：将属于保留类别的标注写入新文件。
• 调用函数：设置输入和输出目录路径以及需要保留的类别索引列表，调用函数执行过滤操作。

确保将input_dir和output_dir变量设置为实际的目录路径，并将classes_to_keep列表设置为你需要保留的类别索引。这个脚本会创建一个新的标注文件，其中只包含所需类别的标注，并将它们保存在指定的输出目录中。