（五）删掉文件夹内所有txt标注文件里面的不需要的类别

import os

class_names = {
    0: 'Pedestrian',
    1: 'Cyclist',
    2: 'Car',
    3: 'Truck',
    4: 'Tram',
    5: 'Tricycle'
}

folder_path = "D:/PycharmProjects/yolov5-master/sodanight/labels/train"

# 遍历文件夹中的每个txt文件
for filename in os.listdir(folder_path):
    if filename.endswith(".txt"):
        file_path = os.path.join(folder_path, filename)

        # 读取txt文件内容
        with open(file_path, 'r') as file:
            lines = file.readlines()

        # 过滤包含需要删除的标注的行
        filtered_lines = []
        for line in lines:
            parts = line.strip().split(" ")
            label = int(parts[0])

            if label not in [4, 5]:  # 排除类别为 4（Tram）和 5（Tricycle）的行
                filtered_lines.append(line)

        # 将过滤后的内容写回txt文件
        with open(file_path, 'w') as file:
            file.writelines(filtered_lines)