Yolov5训练时如何合并COCO数据集与自定义数据集

1. 统一数据集格式 #

YOLOv5采用YOLO格式训练，需先确保两个数据集都适配该格式：

• COCO默认是COCO格式，用YOLOv5自带脚本转成YOLO格式，运行命令：

  python data/coco2yolo.py --coco-dir path/to/coco --output-dir path/to/coco_yolo
  

  转换后会生成images和labels文件夹，标签文件为对应图片的txt，每行格式是class_id x_center y_center width height（数值已归一化）。
• Roboflow导出数据集时直接选择YOLOv5 PyTorch格式，下载后自带train/val/test的images、labels文件夹，以及data.yaml配置文件。

2. 合并数据集文件 #

合并图片与标签目录 #

• 创建根目录combined_dataset，内部新建train、val、test子目录，每个子目录下再建images和labels文件夹。
• 将COCO转换后的train/images、train/labels文件分别复制到combined_dataset/train/images和combined_dataset/train/labels，同理处理COCO的val/test集。
• 将Roboflow自定义数据集的对应目录文件复制到combined_dataset的对应位置，若文件名重复，需加前缀（如custom_）避免覆盖。

统一类别ID #

这是避免类别冲突的核心步骤：

• 查看COCO的80个类别（ID范围0-79），以及Roboflow自定义数据集的类别ID（默认从0开始）。
• 修改自定义数据集的标签文件，将所有类别ID加上COCO最大类别ID+1（比如80），确保两类数据集ID不重叠。可通过Python脚本批量修改：

  import os
  
  label_dir = "path/to/custom_labels"
  offset = 80  # COCO的总类别数
  
  for filename in os.listdir(label_dir):
      if filename.endswith(".txt"):
          with open(os.path.join(label_dir, filename), 'r') as f:
              lines = f.readlines()
          with open(os.path.join(label_dir, filename), 'w') as f:
              for line in lines:
                  parts = line.strip().split()
                  if parts:
                      new_id = int(parts[0]) + offset
                      f.write(f"{new_id} {' '.join(parts[1:])}\n")
  

3. 生成合并后的配置文件 #

新建combined_data.yaml，内容示例如下：

train: ../combined_dataset/train/images
val: ../combined_dataset/val/images
test: ../combined_dataset/test/images

nc: 80 + N  # N为自定义数据集的类别数
names: [
  'person', 'bicycle', 'car', ...,  # COCO的80个类别名称
  'custom_class1', 'custom_class2', ...  # 自定义类别名称，顺序需与修改后的ID对应
]

4. 启动联合训练 #

用YOLOv5的train.py脚本，指定合并后的配置文件：

python train.py --img 640 --batch 16 --epochs 50 --data combined_data.yaml --weights yolov5s.pt

内容的提问来源于stack exchange，提问作者NSA