YOLO自定义数据集实现K折交叉验证——K-Fold Cross Validation

实现K折交叉验证（K-Fold Cross Validation）对于YOLO（You Only Look Once）自定义数据集的目标检测任务可以显著提升模型的可靠性和泛化能力。

这里有一个问题需要注意：无论你是否使用K折交叉验证，测试集都是你无法绕过的。使用 K 折交叉验证是对模型的训练集和验证集进行多次的划分，通过多轮训练，模型的每一轮验证集都会用作一次评估，从而为每一轮的模型性能提供评价，最终再求出平均值，帮助你评估模型的表现并调优参数。测试集通常在 K 折交叉验证后的最后阶段用于最终评估模型的泛化性能。
也就是说K折交叉验证是用来处理训练集和验证集以此来优化模型训练和参数调整过程的，测试集是独立的，需要对最终的模型进行性能评估

详见这篇文章: 目标检测模型性能评估：mAP50、mAP50-95、Precision 和 Recall 及测试集质量的重要性

1. 数据集准备

首先，你需要确保你的数据集符合YOLO的格式，具体来说，每个图像都有相应的标注文件，格式如下：

• 每行包含：class_id center_x center_y width height。
• class_id 是类别的编号，center_x、center_y 是物体中心的归一化坐标，width 和 height 是物体框的归一化宽度和高度。

例如：（东北大学NEU-DET）

XML标注文件

---

TXT标注文件

---

上面的TXTlabel信息即是符合要求的。

如果你的标注文件并不是TXT格式，而是其他格式，那么你需要使用格式转换脚本（可以自己用Python写一个）转换为TXT格式同时满足YOLO标注文件格式要求。下面是我用来将XML转为TXT格式脚本文件（我的数据集是东北大学NEU-DET）

import os
import xml.etree.ElementTree as ET############################################转换核心部分############################################################## 类别映射为类ID
CLASS_MAPPING = {"crazing": 0,  # 类别 "crazing" 对应的 ID"inclusion": 1,  # 添加更多类别"patches": 2,"pitted_surface": 3,"rolled-in_scale": 4,"scratches": 5
}# 转换单个 XML 文件为 YOLO 格式
def convert_xml_to_yolo(xml_file, output_dir):tree = ET.parse(xml_file)root = tree.getroot()# 获取图片尺寸size = root.find('size')img_width = int(size.find('width').text)img_height = int(size.find('height').text)# 输出文件名和路径filename = os.path.splitext(root.find('filename').text)[0]txt_file_path = os.path.join(output_dir, f"{filename}.txt")with open(txt_file_path, 'w') as txt_file:for obj in root.findall('object'):class_name = obj.find('name').textif class_name not in CLASS_MAPPING:print(f"警告：类别 {class_name} 未映射！跳过。")continueclass_id = CLASS_MAPPING[class_name]# 获取边界框坐标bndbox = obj.find('bndbox')xmin = int(bndbox.find('xmin').text)ymin = int(bndbox.find('ymin').text)xmax = int(bndbox.find('xmax').text)ymax = int(bndbox.find('ymax').text)# 转换为YOLO格式x_center = ((xmin + xmax) / 2) / img_widthy_center = ((ymin + ymax) / 2) / img_heightwidth = (xmax - xmin) / img_widthheight = (ymax - ymin) / img_height# 写入TXT文件txt_file.write(f"{class_id} {x_center:.6f} {y_center:.6f} {width:.6f} {height:.6f}\n")print(f"转换完成：{txt_file_path}")############################################批处理部分############################################################## 批量处理 XML 文件
def batch_convert_xml_to_yolo(xml_dir, output_dir):os.makedirs(output_dir, exist_ok=True)xml_files = [f for f in os.listdir(xml_dir) if f.endswith('.xml')]for xml_file in xml_files:xml_path = os.path.join(xml_dir, xml_file)convert_xml_to_yolo(xml_path, output_dir)print(f"全部转换完成！共处理 {len(xml_files)} 个文件。")############################################转换结果写入文件及文件夹######################################################## 输入文件夹路径
xml_dir = r"NEU-DET\ANNOTATIONS"  # XML文件夹路径
output_dir = r"datasets\mydata\labels"  # 输出TXT文件夹路径
batch_convert_xml_to_yolo(xml_dir, output_dir)

现在假设你已经准备好了数据集（例如自定义的水果检测数据集），其中图像和标注文件分别存储在 images 和 labels 目录下。

2. 必要的Python包

你需要安装一些必要的Python库：

pip install -U ultralytics scikit-learn pandas pyyaml

3. 数据集标注和类定义

假设你有一个 data.yaml 文件，它定义了数据集的路径和类别。一个示例 data.yaml 文件可能如下：

train: ./Fruit-detection/images/train
val: ./Fruit-detection/images/val
names:0: Apple1: Grapes2: Pineapple3: Orange4: Banana5: Watermelon

确保你的数据集标注文件（例如 train 和 val 目录中的标注文件）符合此结构。

4. 数据准备和生成特征向量

你需要先生成一个表示数据集的特征向量（每个图像包含每个类的数量）。以下是生成特征向量的代码：

import pandas as pd
from pathlib import Path
from collections import Counter
import yaml# 设置数据集路径
dataset_path = Path("./Fruit-detection")
labels = sorted(dataset_path.rglob("*labels/*.txt"))  # 读取所有标注文件# 读取data.yaml文件，提取类标签
yaml_file = "path/to/data.yaml"
with open(yaml_file, "r", encoding="utf8") as y:classes = yaml.safe_load(y)["names"]# 初始化一个空的DataFrame
cls_idx = sorted(classes.keys())
index = [label.stem for label in labels]  # 使用文件名作为索引
labels_df = pd.DataFrame([], columns=cls_idx, index=index)# 统计每个类的实例数量
for label in labels:lbl_counter = Counter()with open(label, "r") as lf:lines = lf.readlines()for line in lines:lbl_counter[int(line.split(" ")[0])] += 1labels_df.loc[label.stem] = lbl_counterlabels_df = labels_df.fillna(0.0)  # 填充缺失值为0

5. 使用K折交叉验证进行数据拆分

使用 sklearn.model_selection.KFold 来拆分数据集。这里我们使用5折交叉验证（k=5），你可以根据需要调整 k 的值。

from sklearn.model_selection import KFoldksplit = 5
kf = KFold(n_splits=ksplit, shuffle=True, random_state=20)  # 设置随机种子以便结果可复现# 获取数据集的索引拆分
kfolds = list(kf.split(labels_df))# 显示每个fold的训练和验证集
folds_df = pd.DataFrame(index=index, columns=[f"split_{n}" for n in range(1, ksplit + 1)])
for i, (train, val) in enumerate(kfolds, start=1):folds_df[f"split_{i}"].loc[labels_df.iloc[train].index] = "train"folds_df[f"split_{i}"].loc[labels_df.iloc[val].index] = "val"

6. 计算每个fold的标签分布

为了确保每个fold的类别分布平衡，可以计算每个fold中每个类的数量比例。

fold_lbl_distrb = pd.DataFrame(index=[f"split_{n}" for n in range(1, ksplit + 1)], columns=cls_idx)
for n, (train_indices, val_indices) in enumerate(kfolds, start=1):train_totals = labels_df.iloc[train_indices].sum()val_totals = labels_df.iloc[val_indices].sum()# 计算验证集与训练集的标签比例ratio = val_totals / (train_totals + 1e-7)  # 避免除0错误fold_lbl_distrb.loc[f"split_{n}"] = ratio

7. 创建K折数据集文件夹和YAML文件

为每个fold创建训练和验证数据集的文件夹，并生成相应的 dataset.yaml 配置文件。

import shutil
import datetimesave_path = Path(dataset_path / f"{datetime.date.today().isoformat()}_{ksplit}-Fold_Cross-val")
save_path.mkdir(parents=True, exist_ok=True)# 创建目录和YAML文件
ds_yamls = []
for split in folds_df.columns:split_dir = save_path / splitsplit_dir.mkdir(parents=True, exist_ok=True)(split_dir / "train" / "images").mkdir(parents=True, exist_ok=True)(split_dir / "train" / "labels").mkdir(parents=True, exist_ok=True)(split_dir / "val" / "images").mkdir(parents=True, exist_ok=True)(split_dir / "val" / "labels").mkdir(parents=True, exist_ok=True)dataset_yaml = split_dir / f"{split}_dataset.yaml"ds_yamls.append(dataset_yaml)with open(dataset_yaml, "w") as ds_y:yaml.safe_dump({"path": split_dir.as_posix(),"train": "train","val": "val","names": classes,}, ds_y)# 复制图像和标签文件到对应的目录
images = sorted((dataset_path / "images").rglob("*"))
for image, label in zip(images, labels):for split, k_split in folds_df.loc[image.stem].items():img_to_path = save_path / split / k_split / "images"lbl_to_path = save_path / split / k_split / "labels"shutil.copy(image, img_to_path / image.name)shutil.copy(label, lbl_to_path / label.name)

8. 训练YOLO模型

创建一个YOLO模型并使用每个fold的数据进行训练。训练完成后，你可以保存模型并记录性能指标。

from ultralytics import YOLOweights_path = "path/to/weights.pt"  # YOLO预训练权重文件路径
model = YOLO(weights_path, task="detect")# 训练每个fold的数据
results = {}
batch = 16
epochs = 100
project = "kfold_demo"for k in range(ksplit):dataset_yaml = ds_yamls[k]model.train(data=dataset_yaml, epochs=epochs, batch=batch, project=project)results[k] = model.metrics  # 保存训练结果

9. 结果分析

你可以从 results 中提取每个fold的训练指标进行进一步分析。例如，可以计算每个fold的mAP（mean Average Precision）并进行比较，确保模型的稳定性和泛化能力。

---

结论

通过上述步骤，你可以在YOLO自定义数据集上实现K折交叉验证。K折交叉验证的优点是能够减少模型过拟合的风险，确保模型在不同数据划分上的泛化能力，提升其性能可靠性。

这些步骤是通用的，可以根据自己的数据集进行修改和优化。