yolo标签文件根据类别划分到不同的文件夹

import os
import shutil

# 定义类别索引到类别名称的映射
class_mapping = {
    0: 'indicative',
    1: 'prohibitory',
    2: 'warning',
    # 可以继续添加更多的类别映射
}

# 原始文件夹和新文件夹路径
label_source_dir = '/home/zl/yolov5-shiyan（复件）/VOCdevkit/VOCdevkit/labels/val'
output_base_dir = '/home/zl/yolov5-shiyan（复件）/VOCdevkit/VOCdevkit/labels/1'

# 创建输出文件夹的类别子文件夹
for class_name in class_mapping.values():
    os.makedirs(os.path.join(output_base_dir, class_name), exist_ok=True)

# 获取所有标签文件列表
label_files = os.listdir(label_source_dir)

# 遍历标签文件，移动到对应的类别文件夹
for label_file in label_files:
    if label_file.endswith('.txt'):
        label_file_path = os.path.join(label_source_dir, label_file)
        with open(label_file_path, 'r') as f:
            first_line = f.readline().strip()

        # 解析第一行的YOLO格式标签信息
        parts = first_line.split()
        if len(parts) > 0:
            class_index = int(parts[0])
            if class_index in class_mapping:
                class_name = class_mapping[class_index]
                destination_dir = os.path.join(output_base_dir, class_name)
                shutil.move(label_file_path, os.path.join(destination_dir, label_file))
                print(f"Moved {label_file} to {class_name}/")

print("Finished moving label files based on classes.")

注意：一个图像的标签文件可能会有不同的分类，此处是按照txt文件第一行的类别进行划分的