按yolo格式标签分类图像

import os
import shutil

def classify_images_and_labels(images_dir, labels_dir, output_root, class_names=None):
    """
    根据YOLO标签同时分类图片和标签文件
    
    参数:
        images_dir: 原始图片目录
        labels_dir: 原始标签目录
        output_root: 输出根目录
        class_names: 可选，类别名称列表
    """
    # 创建输出目录
    os.makedirs(output_root, exist_ok=True)
    
    # 获取所有标签文件
    label_files = [f for f in os.listdir(labels_dir) if f.endswith('.txt')]
    
    for label_file in label_files:
        # 查找对应的图片文件
        base_name = os.path.splitext(label_file)[0]
        img_path = None
        
        # 尝试常见图片格式
        for ext in ['.jpg', '.jpeg', '.png', '.bmp', '.tif']:
            test_path = os.path.join(images_dir, base_name + ext)
            if os.path.exists(test_path):
                img_path = test_path
                break
        
        if not img_path:
            print(f"注意: 未找到图片 {base_name}")
            continue
        
        # 读取标签文件获取类别
        label_path = os.path.join(labels_dir, label_file)
        try:
            with open(label_path, 'r') as f:
                lines = f.readlines()
        except:
            print(f"错误: 无法读取标签文件 {label_file}")
            continue
        
        # 收集所有类别ID
        class_ids = set()
        for line in lines:
            line = line.strip()
            if line:
                parts = line.split()
                if len(parts) >= 1:
                    try:
                        class_id = int(parts[0])
                        class_ids.add(class_id)
                    except ValueError:
                        continue
        
        # 如果没有有效类别，跳过
        if not class_ids:
            print(f"注意: 标签文件 {label_file} 中没有有效类别")
            continue
        
        # 复制到每个对应的类别文件夹
        for class_id in class_ids:
            # 确定类别文件夹名称
            if class_names and class_id < len(class_names):
                class_folder = class_names[class_id]
            else:
                class_folder = f"class_{class_id}"
            
            # 创建图片和标签的目标目录
            img_dest_dir = os.path.join(output_root, class_folder, "images")
            label_dest_dir = os.path.join(output_root, class_folder, "labels")
            os.makedirs(img_dest_dir, exist_ok=True)
            os.makedirs(label_dest_dir, exist_ok=True)
            
            # 构建目标路径
            img_dest_path = os.path.join(img_dest_dir, os.path.basename(img_path))
            label_dest_path = os.path.join(label_dest_dir, label_file)
            
            # 复制图片文件（如果不存在）
            if not os.path.exists(img_dest_path):
                shutil.copy2(img_path, img_dest_path)
            
            # 复制标签文件（如果不存在）
            if not os.path.exists(label_dest_path):
                shutil.copy2(label_path, label_dest_path)

if __name__ == "__main__":
    # 使用示例
    classify_images_and_labels(
        images_dir="path/to/images",
        labels_dir="path/to/labels",
        output_root="path/to/classified_data",
        class_names=["cat", "dog", "bird"]  # 可选类别名称
    )