YOLOv5以txt或json格式输出预测结果的方法详解

下面我来详细讲解“YOLOv5以txt或json格式输出预测结果的方法详解”。

1. 前提条件

在讲解输出预测结果的方法之前，需要先确保已经完成以下条件：

• 安装好YOLOv5
• 训练好模型
• 准备好测试图片或视频
• 确定输出格式(txt或json)

以下步骤将以txt格式为例进行讲解。

2. 如何输出txt格式的预测结果

2.1 使用--save-txt参数输出预测结果

使用YOLOv5的命令行工具进行预测时，可以使用--save-txt参数来输出预测结果。具体操作如下所示：

python detect.py --weights path/to/weights.pt --img 640 --conf 0.25 --source path/to/source --save-txt

其中：

• --weights参数指定了训练好的模型权重文件的路径。
• --img参数指定了输入图片的大小，单位为像素。
• --conf参数指定了置信度的阈值，低于这个阈值的预测框将被过滤掉。
• --source参数指定了待预测的图片或视频文件的路径。
• --save-txt参数则是输出txt格式的预测结果。

运行以上命令后，程序会在同级目录下生成一个与输入图片同名的txt文件，文件内容格式如下所示：

0 0.32 265 535 545 768
0 0.22 295 776 388 1008
0 0.13 701 546 819 822

每行表示一个预测框，格式为：

[class_id] [confidence] [x_center] [y_center] [width] [height]

其中，[class_id]表示类别编号（从0开始），[confidence]表示预测框的置信度，[x_center]和[y_center]表示预测框的中心点坐标（相对于图片大小），[width]和[height]表示预测框的宽度和高度（也是相对于图片大小的比例）。

2.2 自定义txt格式的输出结果

此外，如果需要更加详细或特殊的预测结果格式，则可以自定义txt格式的输出结果。具体方法是在YOLOv5的代码中对save_txt.py文件进行修改。

首先，找到save_txt.py文件，该文件的默认路径为yolov5/utils/save_txt.py。然后，打开文件，可以看到其中的代码如下所示：

import os


def results_to_txt(results, save_dir):
    for img_i, (path, preds) in enumerate(results):
        txt_path = os.path.join(save_dir, os.path.basename(path).replace('.jpg', '.txt'))
        with open(txt_path, 'w') as file:
            for *xyxy, conf, cls in preds:
                file.write(('%g ' * 6 + '\n') % (*cls, conf, *xyxy))

在该文件中，results_to_txt(results, save_dir)函数负责将预测结果写入txt文件，具体格式为：

[class_id] [confidence] [x_min] [y_min] [x_max] [y_max]

如果需要自定义输出结果格式，可以按照自己需求修改此函数中的代码。例如，以下代码实现了在每个预测框前增加一行标注信息：

import os


def results_to_txt(results, save_dir):
    for img_i, (path, preds) in enumerate(results):
        txt_path = os.path.join(save_dir, os.path.basename(path).replace('.jpg', '.txt'))
        with open(txt_path, 'w') as file:
            file.write('image_path: %s\n' % path)
            for *xyxy, conf, cls in preds:
                file.write(('%g ' * 6 + '\n') % (*cls, conf, *xyxy))

运行以上代码后，输出的预测结果文件将包含一个标注信息行，格式为：

image_path: [path/to/image]
[class_id] [confidence] [x_min] [y_min] [x_max] [y_max]

3. 如何输出json格式的预测结果

要输出json格式的预测结果，可以使用YOLOv5的--save-json参数。下面是具体操作步骤：

python detect.py --weights path/to/weights.pt --img 640 --conf 0.25 --source path/to/source --save-json

其中，--source和--save-txt参数与上文相同，这里不再赘述。其他参数含义如下：

• --save-json表示将预测结果保存为json格式。
• --entity-format指定实体输出格式，有三种选项：str（字符串）、xyxy（左上角和右下角坐标）和xywh（中心点坐标和宽高）。
• --json-pretty表示在输出的json文件中将预测框等信息进行格式化，方便查看。

运行以上命令后，程序会在同级目录下生成一个与输入图片同名的json文件，文件内容格式如下所示：

{
    "file": "path/to/image.jpg",
    "image_size": [
        1280,
        720
    ],
    "objects": [
        {
            "label": "person",
            "confidence": 0.582,
            "bbox": [
                448,
                90,
                1079,
                700
            ]
        },
        {
            "label": "bicycle",
            "confidence": 0.112,
            "bbox": [
                773,
                230,
                1243,
                659
            ]
        }
    ]
}

其中，"file"为输出图片的路径；"image_size"为输出图片的大小；"objects"为预测结果列表，包括每个预测框的类别、置信度和坐标。

4. 总结

以上就是YOLOv5以txt或json格式输出预测结果的方法详解。通过以上方法，我们可以方便地输出不同格式的预测结果，并可以根据需要自定义输出结果格式。