YOLOv5数据集划分与JSON生成全流程：从小白到APs指标计算

YOLOv5实战：从数据集划分到APs指标计算的完整指南

如果你刚接触目标检测，面对一堆图片和标签文件，可能会感到无从下手。我刚开始用YOLOv5时，光是搞清楚数据集该怎么组织、如何评估模型性能就花了不少时间。特别是那个APs指标——听起来很专业，实际操作起来却需要一系列看似琐碎但至关重要的步骤。今天我就把自己踩过的坑、总结的经验完整分享出来，帮你跳过那些让人“抑郁”的调试过程，直接掌握从原始数据到最终评估的全流程。

这篇文章面向的是已经准备好YOLO格式数据集，但不知道如何系统地进行数据划分、格式转换，并最终计算出APs等关键指标的开发者。我们会从最基础的文件结构开始，一步步讲解每个环节的操作细节和背后的原理。你会发现，只要理解了每个步骤的意义，整个过程其实并不复杂。

1. 理解YOLO数据集的基础结构

在开始任何操作之前，我们需要先搞清楚YOLO格式数据集到底长什么样。很多人直接下载了数据集就开始跑训练，结果遇到各种路径错误、格式问题，根本原因是对数据组织方式理解不够深入。

1.1 YOLO格式的核心要素

YOLO格式的数据集通常包含两个核心部分：图片文件和对应的标签文件。图片可以是常见的格式如JPG、PNG等，而标签文件则是纯文本的TXT文件，每个文件与图片同名（仅扩展名不同）。

一个典型的YOLO标签文件内容如下：

0 0.45 0.32 0.12 0.08
1 0.67 0.54 0.09 0.15

每行代表一个目标对象，包含5个数值：

• 第一个数字：类别ID（从0开始）
• 后面四个数字：归一化的边界框坐标（x_center, y_center, width, height）

这里的归一化指的是相对于图片宽高的比例值。例如，x_center=0.45表示边界框中心位于图片宽度的45%位置。

1.2 文件命名与组织规范

在实际项目中，我强烈建议采用统一的命名规则。虽然YOLO本身不强制要求，但良好的命名习惯能避免很多不必要的麻烦：

# 推荐的命名模式示例
图片文件: 000001.jpg, 000002.jpg, 000003.jpg, ...
标签文件: 000001.txt, 000002.txt, 000003.txt, ...

注意：避免在文件名中使用特殊字符和空格，这可能导致某些脚本解析失败。纯数字命名虽然看起来简单，但在处理大量文件时确实能减少排序和匹配的问题。

关于文件夹结构，YOLOv5期望的默认结构是这样的：

datasets/
├── images/
│   ├── train/      # 训练集图片
│   └── val/        # 验证集图片
└── labels/
    ├── train/      # 训练集标签
    └── val/        # 验证集标签

但很多时候我们拿到的是未划分的原始数据，这就需要我们手动进行数据集划分。

2. 数据集划分：科学分配训练与验证样本

数据集划分是机器学习项目中的关键一步，划分不合理可能导致模型过拟合或评估结果不可靠。对于目标检测任务，特别是小目标检测，划分时还需要考虑类别平衡问题。

2.1 划分比例的选择策略

常见的划分比例有7:3、8:2等，但具体选择需要考虑数据集大小：

对于小目标检测，我建议验证集比例不低于20%，因为小目标的检测性能波动较大，需要足够的验证样本来准确评估。

2.2 随机划分与分层抽样的实现

简单的随机划分可能导致某些类别在验证集中样本过少。下面是一个改进的划分脚本，它确保了每个类别在训练集和验证集中都有代表性样本：

import os
import random
import shutil
from collections import defaultdict

def stratified_split(img_dir, label_dir, output_dir, train_ratio=0.8):
    """
    分层抽样划分数据集，确保类别分布均衡
    
    参数:
        img_dir: 原始图片目录
        label_dir: 原始标签目录
        output_dir: 输出目录
        train_ratio: 训练集比例
    """
    # 创建输出目录结构
    os.makedirs(os.path.join(output_dir, 'images', 'train'), exist_ok=True)
    os.makedirs(os.path.join(output_dir, 'images', 'val'), exist_ok=True)
    os.makedirs(os.path.join(output_dir, 'labels', 'train'), exist_ok=True)
    os.makedirs(os.path.join(output_dir, 'labels', 'val'), exist_ok=True)
    
    # 按类别组织文件
    class_files = defaultdict(list)
    
    # 遍历所有图片文件
    img_files = [f for f in os.listdir(img_dir) 
                 if f.lower().endswith(('.jpg', '.jpeg', '.png', '.bmp'))]
    
    for img_file in img_files:
        label_file = os.path.splitext(img_file)[0] + '.txt'
        label_path = os.path.join(label_dir, label_file)
        
        if not os.path.exists(label_path):
            continue
            
        # 读取标签文件，确定主要类别
        with open(label_path, 'r') as f:
            lines = f.readlines()
            if lines:
                # 取第一个目标的类别作为该图片的代表类别
                # 对于多类别图片，可以更复杂的策略
                first_class = int(lines[0].split()[0])
                class_files[first_class].append((img_file, label_file))
    
    # 对每个类别分别划分
    train_files = []
    val_files = []
    
    for class_id, files in class_files.items():
        random.shuffle(files)
        split_idx = int(len(files) * train_ratio)
        
        train_files.extend(files[:split_idx])
        val_files.extend(files[split_idx:])
    
    # 打乱最终的文件列表
    random.shuffle(train_files)
    random.shuffle(val_files)
    
    # 复制文件到对应目录
    def copy_files(file_list, split_type):
        for img_file, label_file in file_list:
            # 复制图片
            src_img = os.path.join(img_dir, img_file)
            dst_img = os.path.join(output_dir, 'images', split_type, img_file)
            shutil.copy2(src_img, dst_img)
            
            # 复制标签
            src_label = os.path.join(label_dir, label_file)
            dst_label = os.path.join(output_dir, 'labels', split_type, label_file)
            shutil.copy2(src_label, dst_label)
    
    copy_files(train_files, 'train')
    copy_files(val_files, 'val')
    
    print(f"划分完成: 训练集 {len(train_files)} 张, 验证集 {len(val_files)} 张")
    return len(train_files),