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

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

1. 数据集准备

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

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

假设你已经准备好了数据集（例如自定义的水果检测数据集），其中图像和标注文件分别存储在 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折交叉验证的优点是能够减少模型过拟合的风险，确保模型在不同数据划分上的泛化能力，提升其性能可靠性。

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