如何使用自定义训练的Mask RCNN自动生成缺陷图像的标注JSON文件

核心实现步骤 #

• 步骤1：批量推理未标注图像，提取模型输出字段
  你已经在detectron2中完成了自定义Mask RCNN模型的训练，直接基于现有推理接口加载训练好的权重，对剩余410张未标注图像执行推理：
  调用predictor(输入图像)得到输出的Instances对象后，提取以下核心字段：
  • pred_classes：预测的缺陷类别ID
  • pred_boxes：实例边界框坐标
  • pred_masks：实例分割二值掩码
  • scores：预测置信度
    你可以根据验证集的精度表现设置置信度过滤阈值，通常取验证集AP对应的置信度值下浮0.1作为阈值，过滤掉置信度过低的错误预测结果。
• 步骤2：推理结果转LabelMe兼容JSON格式
  单张图像对应生成单份LabelMe格式JSON的映射规则如下：
  1. 对每个保留的实例掩码，调用cv2.findContours()提取轮廓坐标，转换为LabelMe要求的多边形点集格式
  2. label字段填入缺陷类别ID映射后的类别名称
  3. 可额外新增自定义字段存储预测置信度，方便后续人工校验时优先排查低置信度标注
  4. 保持生成的JSON文件名与对应图像文件名完全匹配，存储路径和图像路径的相对关系和你之前人工标注的数据集保持一致，避免LabelMe加载失败。
• 步骤3：人工轻量校验修正标注结果
  自动生成的标注在用于后续训练前必须做轻量人工校验，你直接用LabelMe批量加载生成的标注文件，仅需要完成以下操作即可：
  • 优先校验置信度低于阈值的标注结果，删除误检实例、补充漏检实例、修正错误的类别标注
  • 高置信度的标注结果快速过检即可，这类预测准确率普遍较高
    按照这个流程下来，标注效率可以提升3~6倍，大幅降低人工标注耗时。

可选优化方案 #

• 完成本次所有图像标注完成后，将新增的标注数据和原始40张标注数据合并，重新微调训练现有Mask RCNN模型，模型精度会进一步提升，如果后续还有标注需求，自动标注效果会更好，形成数据迭代的正向循环。
• 若你的数据集中存在极小尺寸缺陷，可以针对小目标单独调整置信度阈值，避免漏检。

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