如何修改Detectron2代码将预测输出保存为JSON用于后续标注？

要将Detectron2的预测结果保存为JSON，需要先把模型输出的Instances对象转换为可JSON序列化的结构，再统一写入文件。以下是修改后的完整代码：

import json
import glob
import cv2
from detectron2.utils.visualizer import ColorMode
from google.colab.patches import cv2_imshow  # Colab环境需保留，本地环境可替换为cv2.imshow

# 存储所有图片的预测结果
all_results = []
path = '/content/drive/MyDrive/Detectron2/test/*.png'

for file in glob.glob(path):
    print(file)
    im = cv2.imread(file, 1)
    outputs = predictor(im)
    instances = outputs["instances"].to("cpu")
    
    # 整理当前图片的预测信息
    img_result = {
        "image_path": file,
        "predictions": []
    }
    
    # 提取实例的核心数据
    pred_classes = instances.pred_classes.numpy().tolist()
    pred_boxes = instances.pred_boxes.tensor.numpy().tolist()
    scores = instances.scores.numpy().tolist()
    
    for cls_id, bbox, score in zip(pred_classes, pred_boxes, scores):
        # 转换为可读的类别名称
        cls_name = train_metadata.thing_classes[cls_id]
        img_result["predictions"].append({
            "category_id": cls_id,
            "category_name": cls_name,
            "bbox": [round(coord, 2) for coord in bbox],  # 保留两位小数优化可读性
            "confidence": round(score, 4)
        })
    
    all_results.append(img_result)
    
    # 保留原有的可视化逻辑
    v = Visualizer(im[:, :, ::-1],
                   metadata=train_metadata, 
                   scale=0.5,
                   instance_mode=ColorMode.SEGMENTATION
    )
    out = v.draw_instance_predictions(instances)
    cv2_imshow(out.get_image()[:, :, ::-1])

# 将结果写入JSON文件
with open('/content/drive/MyDrive/Detectron2/predictions.json', 'w') as f:
    json.dump(all_results, f, indent=4)
print("预测结果已保存为predictions.json")

关键说明： #

• 导入json模块实现数据序列化
• 用all_results列表统一存储所有图片的预测数据，每条数据包含图片路径和对应实例的预测详情
• 通过.numpy().tolist()将Detectron2的张量格式转换为JSON支持的列表格式
• 把类别ID转换为对应名称，方便后续标注时直接查看
• indent=4让生成的JSON文件格式更规整，便于人工查看和编辑

内容的提问来源于stack exchange，提问作者Lim Cheng Wei