【YOLO系列PR、F1绘图】更改v5、v7、v8（附v8训练、验证方式），实现调用val.py或者test.py后生成pr.csv，然后再整合绘制到一张图上（使用matplotlib绘制）

原创

已于 2023-06-12 16:02:45 修改 · 2.7w 阅读

680

标签

#matplotlib #YOLO #python

于 2023-04-24 11:23:49 首次发布

文章介绍了如何修改YOLOv7、YOLOv5(v6.1)和YOLOv8的代码，以生成PR_curve.csv和F1_curve.csv文件，这些文件包含了模型验证时的精度和召回率数据。通过调整utils/metrics.py中的函数，可以在验证过程中保存这些关键指标，并提供了绘制PR曲线和F1曲线的方法，以便于对比不同YOLO变体的性能。

1. 前提 + 效果图

不错的链接：YOLOV7训练模型分析
关于map的绘图、loss绘图，可参考：【YOLO系列result中的map、loss、pr绘图】根据v5、v8、v7训练后生成的result文件用matplotlib进行绘图
v5、v8调用val.py，v7调用test.py（作用都是一样的，都是用已训练好权重对测试集进行验证，然后打印出一系列指标）
实现效果：就是将运行val.py/test.py后生成的PR_curve.png中最粗的蓝线整合到同一张图中（注意：本代码最重要的作用是将验证时得到的一系列P、R值给提出来，所以绘图就比较潦草，直接用的matplotlib画的，如果要用于论文中的绘图，一般使用origin）

在这里插入图片描述

同理，可以实现F1_curve.png绘图

在这里插入图片描述

2. 更改步骤

2.1 得到PR_curve.csv和F1_curve.csv

2.1.1 YOLOv7的更改

2.1.1.1 得到PR_curve.csv

在utils/metrics.py中，按住Ctrl+F搜索def plot_pr_curve定位过去，然后如图做更改：

在这里插入图片描述

# Plots ----------------------------------------------------------------------------------------------------------------

def plot_pr_curve(px, py, ap, save_dir='pr_curve.png', names=()):
    # Precision-recall curve
    fig, ax = plt.subplots(1, 1, figsize=(9, 6), tight_layout=True)
    py = np.stack(py, axis=1)

    # lwd edit: 将结果保存在csv中
    pr_dict = dict()    # lwd edit
    pr_dict['px'] = px.tolist() # lwd edit

    if 0 < len(names) < 21:  # display per-class legend if < 21 classes
        for i, y in enumerate(py.T):
            ax.plot(px, y, linewidth=1, label=f'{
     
     names[i]} {
     
     ap[i, 0]:.3f}')  # plot(recall, precision)
            pr_dict[names[i]] = y.tolist()  # lwd edit
    else:
        ax.plot(px, py, linewidth=1, color='grey')  # plot(recall, precision)

    ax.plot(px, py.mean(1), linewidth=3, color='blue', label='all classes %.3f mAP@0.5' % ap[:, 0].mean())

    # ------------------- lwd edit ---------------------- #
    pr_dict['all']