如何在Detectron2中合并自定义训练数据集与其他预训练数据集？

1. 统一数据集格式 #

Detectron2对COCO格式的JSON标注支持最完善，首先将两个数据集统一为COCO格式：

• 如果其中一个是VOC等其他格式，可通过脚本转换为COCO JSON（比如用detectron2/data/datasets/pascal_voc.py中的转换逻辑）。
• 重点对齐类别体系：若两个数据集有重复类别，确保类别名称、ID完全一致；若有新类别，将其加入统一的类别列表。

2. 合并COCO标注文件 #

编写Python脚本合并两个数据集的images、annotations、categories三个核心字段，处理ID冲突问题：

import json

# 加载两个数据集的JSON文件
with open("custom_train.json", "r") as f:
    custom_data = json.load(f)
with open("pretrained_dataset.json", "r") as f:
    pretrained_data = json.load(f)

# 1. 合并并统一类别
category_map = {}
merged_categories = []
# 先导入自定义数据集类别
for cat in custom_data["categories"]:
    if cat["name"] not in category_map:
        category_map[cat["name"]] = len(category_map) + 1  # COCO类别ID从1开始
        merged_categories.append({"id": category_map[cat["name"]], "name": cat["name"]})
# 导入预训练数据集的新类别
for cat in pretrained_data["categories"]:
    if cat["name"] not in category_map:
        category_map[cat["name"]] = len(category_map) + 1
        merged_categories.append({"id": category_map[cat["name"]], "name": cat["name"]})

# 2. 合并图片，处理重复image_id
image_id_map = {}
merged_images = []
current_img_id = 1
# 处理自定义数据集图片
for img in custom_data["images"]:
    image_id_map[img["id"]] = current_img_id
    new_img = img.copy()
    new_img["id"] = current_img_id
    merged_images.append(new_img)
    current_img_id += 1
# 处理预训练数据集图片
for img in pretrained_data["images"]:
    image_id_map[img["id"]] = current_img_id
    new_img = img.copy()
    new_img["id"] = current_img_id
    merged_images.append(new_img)
    current_img_id += 1

# 3. 合并标注，调整image_id和category_id
merged_annotations = []
# 处理自定义数据集标注
for ann in custom_data["annotations"]:
    new_ann = ann.copy()
    new_ann["image_id"] = image_id_map[ann["image_id"]]
    # 映射类别ID
    cat_name = next(c["name"] for c in custom_data["categories"] if c["id"] == ann["category_id"])
    new_ann["category_id"] = category_map[cat_name]
    merged_annotations.append(new_ann)
# 处理预训练数据集标注
for ann in pretrained_data["annotations"]:
    new_ann = ann.copy()
    new_ann["image_id"] = image_id_map[ann["image_id"]]
    cat_name = next(c["name"] for c in pretrained_data["categories"] if c["id"] == ann["category_id"])
    new_ann["category_id"] = category_map[cat_name]
    merged_annotations.append(new_ann)

# 保存合并后的数据集
merged_data = {
    "images": merged_images,
    "annotations": merged_annotations,
    "categories": merged_categories
}
with open("merged_dataset.json", "w") as f:
    json.dump(merged_data, f, indent=2)

3. 注册合并后的数据集 #

在Detectron2中注册合并后的数据集，确保图片路径正确：

from detectron2.data.datasets import register_coco_instances
from detectron2.data import MetadataCatalog, DatasetCatalog

# 注册训练集
register_coco_instances(
    "merged_dataset_train",
    {},
    "path/to/merged_dataset.json",
    "path/to/all_images_dir"  # 若两个数据集图片在不同目录，需先合并到同一目录
)
# 获取元数据
merged_metadata = MetadataCatalog.get("merged_dataset_train")

4. 基于合并数据集训练/微调 #

修改Detectron2配置，指向合并后的数据集并调整参数：

from detectron2.config import get_cfg
import detectron2.model_zoo as model_zoo

cfg = get_cfg()
cfg.merge_from_file(model_zoo.get_config_file("COCO-Detection/faster_rcnn_R_50_FPN_3x.yaml"))

# 替换为合并后的数据集
cfg.DATASETS.TRAIN = ("merged_dataset_train",)
cfg.DATASETS.TEST = ()

# 调整类别数为合并后的总数
cfg.MODEL.ROI_HEADS.NUM_CLASSES = len(merged_categories)

# 加载权重：如果是基于自定义训练模型微调，用自定义权重；如果从头训练，用预训练权重
cfg.MODEL.WEIGHTS = "path/to/your_custom_trained_model.pth"  # 或 model_zoo.get_checkpoint_url(...)

# 其他训练参数调整（批次大小、学习率等）
cfg.SOLVER.IMS_PER_BATCH = 4
cfg.SOLVER.BASE_LR = 0.00025

关键注意事项 #

• 类别名称统一：若两个数据集对同一类别的命名不同（比如"dog"和"canine"），需手动修改为一致名称，否则会被识别为不同类别。
• 数据分布平衡：若两个数据集样本量差异大，可在训练时设置加权采样，避免模型偏向样本多的数据集。
• 图片路径校验：确保所有图片的file_name字段指向正确路径，Detectron2能正常读取。

内容的提问来源于stack exchange，提问作者Rehan Abu Hashir