yolov5模型预测结果保存为json

要将 YOLOv5 模型的预测结果保存为 json 文件，可以按照以下步骤进行操作：

1. 导入必要的库和函数

import json
import torch
from pathlib import Path
from models.experimental import attempt_load
from utils.general import non_max_suppression, scale_coords
from utils.datasets import LoadImages
from utils.torch_utils import select_device

2. 加载模型和设备

weights = 'path/to/weights.pt'
device = select_device('0')
model = attempt_load(weights, map_location=device)
img_size = model.img_size

3. 加载图像并进行预测

# 加载图像
dataset = LoadImages('path/to/images/', img_size=img_size)

# 进行预测
results = []
for path, img, im0s, _ in dataset:
    img = torch.from_numpy(img).to(device)
    img = img.float() / 255.0
    if img.ndimension() == 3:
        img = img.unsqueeze(0)
    pred = model(img, augment=False)[0]
    pred = non_max_suppression(pred, 0.4, 0.5)
    for i, det in enumerate(pred):
        if det is not None and len(det):
            det[:, :4] = scale_coords(img.shape[2:], det[:, :4], im0s.shape).round()
            for *xyxy, conf, cls in det:
                results.append({
                    'image_id': Path(path).stem,
                    'category_id': int(cls),
                    'bbox': xyxy,
                    'score': float(conf)
                })

4. 将预测结果保存为 json 文件

with open('path/to/results.json', 'w') as f:
    json.dump(results, f)

这样，预测结果就会保存在指定的 json 文件中。注意，在保存 bbox 值时，我们使用了 scale_coords 函数将预测框的坐标从图像坐标系转换为原始图像坐标系。