如何让训练好的自定义YOLOv5模型返回含类别、分数及边界框的JSON？

问题描述 #

已完成自定义YOLOv5数据集训练，需求是输入图片后返回包含识别类别、置信度分数、边界框的JSON数据，但运行官方detect.py仅能生成标注后的图片，无法直接输出目标格式的结果。

解决方案 #

有两种可行方式实现需求：

方式一：修改detect.py脚本

在官方detect.py的检测结果处理逻辑（通常在results = model(...)代码段之后），添加结果转JSON的代码，直接输出或保存结构化数据：

import json

# 遍历每张图片的检测结果
for i, det in enumerate(results.pred):
    if len(det):
        img_path = results.files[i]
        detections = []
        # 解析每个检测框的信息
        for *xyxy, conf, cls in det:
            # 转换边界框为浮点型（xyxy格式，可按需改为xywh等）
            bbox = [float(coord) for coord in xyxy]
            # 映射类别ID到类别名称
            cls_name = model.names[int(cls)]
            detections.append({
                "class": cls_name,
                "confidence": float(conf),
                "bbox": bbox
            })
        # 打印JSON结果
        print(json.dumps({"image_path": img_path, "detections": detections}, indent=2))
        # 可选：保存结果到JSON文件
        # with open(f"{os.path.splitext(img_path)[0]}.json", "w") as f:
        #     json.dump({"image_path": img_path, "detections": detections}, f, indent=2)

修改后运行detect.py，即可直接得到JSON格式的检测输出。

方式二：使用独立脚本（优化版）

你已实现的独立脚本方案更灵活，适合集成到其他调用程序中，以下是针对性优化建议及优化后代码：

1. 用pathlib处理路径，避免跨系统分隔符问题
2. 直接映射类别ID到名称，无需手动对应
3. 结构化结果并直接输出JSON，方便调用方解析
4. 支持批量图片处理

优化后的脚本：

import json
import torch
from pathlib import Path

# 自动获取项目根目录（鲁棒性更强）
ROOT = Path(__file__).resolve().parent

# 加载本地训练好的YOLOv5模型
model = torch.hub.load(
    str(ROOT / "yolov5"),
    "custom",
    path=str(ROOT / "yolov5/runs/train/lemon3/weights/best.pt"),
    source="local",
    force_reload=True
)

# 支持单张/多张图片输入
image_paths = [ROOT / "dataset/test/images/fridge2.jpg"]
# 批量处理示例：image_paths = [ROOT / "dataset/test/images/img1.jpg", ROOT / "dataset/test/images/img2.jpg"]

output_results = []
for img_path in image_paths:
    results = model(str(img_path))
    # 提取归一化坐标的检测结果（需原始像素坐标可改用results.xyxy[0].numpy()）
    det_list = results.xyxyn[0].numpy()
    
    detections = []
    for det in det_list:
        x1, y1, x2, y2, conf, cls_id = det
        detections.append({
            "class": model.names[int(cls_id)],
            "confidence": round(float(conf), 4),  # 保留4位小数优化输出
            "bbox": {
                "x1": float(x1),
                "y1": float(y1),
                "x2": float(x2),
                "y2": float(y2),
                "normalized": True  # 标记坐标是否为归一化值
            }
        })
    
    output_results.append({
        "image_path": str(img_path),
        "detections": detections
    })

# 输出JSON格式结果
print(json.dumps(output_results, indent=2, ensure_ascii=False))
# 可选：保存结果到文件
# with open(ROOT / "detection_results.json", "w", encoding="utf-8") as f:
#     json.dump(output_results, f, indent=2, ensure_ascii=False)

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