YOLOv5怎么以txt或json格式输出预测结果

YOLOv5 是一个流行的目标检测模型，广泛应用于计算机视觉任务中。在实际应用中，我们通常需要将模型的预测结果以特定的格式输出，以便后续处理或分析。本文将详细介绍如何使用 YOLOv5 以 txt 或 json 格式输出预测结果。

1. YOLOv5 预测结果的基本格式

YOLOv5 的预测结果通常包含以下信息：

• 目标类别（class）
• 目标边界框（bounding box）
• 目标置信度（confidence）

这些信息可以通过 YOLOv5 的推理脚本直接输出到控制台，或者保存到文件中。

2. 以 txt 格式输出预测结果

2.1 修改 YOLOv5 的推理脚本

YOLOv5 的推理脚本 detect.py 默认会将预测结果保存为图像文件，并在图像上绘制检测框。为了将预测结果保存为 txt 文件，我们需要对 detect.py 进行一些修改。

首先，找到 detect.py 中保存结果的部分代码：

# Save results (image with detections)
if save_img:
    if dataset.mode == 'image':
        cv2.imwrite(save_path, im0)
    else:  # 'video' or 'stream'
        if vid_path != save_path:  # new video
            vid_path = save_path
            if isinstance(vid_writer, cv2.VideoWriter):
                vid_writer.release()  # release previous video writer
            fps = vid_cap.get(cv2.CAP_PROP_FPS)
            w = int(vid_cap.get(cv2.CAP_PROP_FRAME_WIDTH))
            h = int(vid_cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
            vid_writer = cv2.VideoWriter(save_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (w, h))
        vid_writer.write(im0)

在这段代码之后，我们可以添加保存 txt 文件的逻辑。假设我们希望每个图像对应一个 txt 文件，文件中每一行表示一个检测到的目标，格式为 class x_center y_center width height confidence。

# Save results to txt file
if save_txt:
    txt_path = os.path.splitext(save_path)[0] + '.txt'
    with open(txt_path, 'w') as f:
        for *xyxy, conf, cls in reversed(det):
            x_center = (xyxy[0] + xyxy[2]) / 2 / im0.shape[1]
            y_center = (xyxy[1] + xyxy[3]) / 2 / im0.shape[0]
            width = (xyxy[2] - xyxy[0]) / im0.shape[1]
            height = (xyxy[3] - xyxy[1]) / im0.shape[0]
            f.write(f'{int(cls)} {x_center:.6f} {y_center:.6f} {width:.6f} {height:.6f} {conf:.6f}\n')

2.2 运行推理脚本

在修改完 detect.py 后，我们可以通过以下命令运行推理脚本，并指定保存 txt 文件的路径：

python detect.py --source path/to/images --save-txt --output path/to/output

其中，--source 指定输入图像或视频的路径，--save-txt 表示保存 txt 文件，--output 指定输出路径。

3. 以 json 格式输出预测结果

3.1 修改 YOLOv5 的推理脚本

与 txt 格式类似，我们可以通过修改 detect.py 来保存 json 格式的预测结果。json 格式通常更适合存储结构化数据，并且可以包含更多的信息。

在 detect.py 中，找到保存结果的部分代码，并添加保存 json 文件的逻辑：

import json

# Save results to json file
if save_json:
    json_path = os.path.splitext(save_path)[0] + '.json'
    results = []
    for *xyxy, conf, cls in reversed(det):
        x_center = (xyxy[0] + xyxy[2]) / 2 / im0.shape[1]
        y_center = (xyxy[1] + xyxy[3]) / 2 / im0.shape[0]
        width = (xyxy[2] - xyxy[0]) / im0.shape[1]
        height = (xyxy[3] - xyxy[1]) / im0.shape[0]
        result = {
            'class': int(cls),
            'x_center': float(x_center),
            'y_center': float(y_center),
            'width': float(width),
            'height': float(height),
            'confidence': float(conf)
        }
        results.append(result)
    with open(json_path, 'w') as f:
        json.dump(results, f, indent=4)

3.2 运行推理脚本

在修改完 detect.py 后，我们可以通过以下命令运行推理脚本，并指定保存 json 文件的路径：

python detect.py --source path/to/images --save-json --output path/to/output

其中，--source 指定输入图像或视频的路径，--save-json 表示保存 json 文件，--output 指定输出路径。

4. 总结

通过修改 YOLOv5 的推理脚本 detect.py，我们可以轻松地将预测结果保存为 txt 或 json 格式。txt 格式适合简单的数据存储，而 json 格式则更适合存储结构化数据。根据实际需求选择合适的输出格式，可以大大提高后续数据处理的效率。

希望本文对你理解和使用 YOLOv5 有所帮助。如果你有任何问题或建议，欢迎在评论区留言讨论。