如何用COCO格式数据集与Detectron2训练无骨架自定义关键点检测

一、自定义COCO格式数据集调整 #

Detectron2对COCO格式的关键点检测兼容度很高，哪怕是单关键点、无骨架的场景，只需要针对性修改几个核心字段即可：

• 标注阶段：针对每张图里的汽车、自行车实例，仅标注一个车轮关键点，记录其坐标（x,y）和可见性（直接填2表示完全可见即可）。
• JSON文件字段修改：
  • categories数组：每个类别（汽车、自行车）需添加keypoints和skeleton字段。keypoints只放一个元素（比如["wheel"]），skeleton设为空数组[]即可，示例：

    "categories": [
        {"id": 1, "name": "car", "supercategory": "vehicle", "keypoints": ["wheel"], "skeleton": []},
        {"id": 2, "name": "bike", "supercategory": "vehicle", "keypoints": ["wheel"], "skeleton": []}
    ]
    

  • annotations数组：每个实例的keypoints字段是长度为3的数组[x, y, visibility]，同时num_keypoints必须设为1，示例：

    "annotations": [
        {"id": 1, "image_id": 1, "category_id": 1, "bbox": [100, 200, 300, 400], "area": 120000, "iscrowd": 0, "keypoints": [150, 300, 2], "num_keypoints": 1}
    ]
    

二、Detectron2训练配置修改 #

基于你之前的实例分割经验，只需要调整关键点检测相关的配置项：

• 注册自定义数据集：用register_coco_instances直接注册你的数据集，无需额外处理关键点字段，Detectron2会自动识别：

  from detectron2.data.datasets import register_coco_instances
  register_coco_instances("vehicle_kp_train", {}, "./train_kp.json", "./train_images")
  register_coco_instances("vehicle_kp_val", {}, "./val_kp.json", "./val_images")
  

• 修改模型核心参数：
  1. 选择KeypointRCNN作为基础模型，加载COCO预训练权重。
  2. 修改MODEL.ROI_KEYPOINT_HEAD.NUM_KEYPOINTS为1，匹配单关键点场景。
  3. 修改MODEL.ROI_HEADS.NUM_CLASSES为2（对应汽车、自行车两个类别）。
  4. 若不需要关键点骨架相关的损失计算，Detectron2会自动忽略空的skeleton字段，无需额外关闭。

三、其他可行方案 #

如果不想用COCO格式，也可以自定义数据集加载器，但COCO格式是Detectron2官方支持最好的，无需额外写代码，优先推荐。另外，若你的需求只是定位车轮，也可以把车轮当成独立的目标检测类别，但这样丢失了车轮与父实例（汽车/自行车）的关联，不如关键点检测方案精准。

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