yolov5、yolov8训练结果绘制

将模型训练的结果绘制在一个图上

import os

import matplotlib.pyplot as plt
import pandas as pd

if __name__ == '__main__':
    # 模型路径
    root_path = r'D:\lg\BaiduSyncdisk\project\person_code\yolov8\runs\train\安全帽'
    # 图保存路径
    output_path = r'./output/plot_results_csv'
    # 需要绘制的模型
    lst_results_csv = ['yolov8s', 'yolov8s-attention-BAM', 'yolov8s-attention-CA', 'yolov8s-attention-CBAM', 'yolov8s-attention-ECA',
                       'yolov8s-attention-SE',
                       # 主干
                       'yolov8s-backbone-efficientnet_b0', 'yolov8s-backbone-mobilenetv3_small_050',
                       'yolov8s-backbone-mobilenetv3_small_100', 'yolov8s-backbone-renset18', 'yolov8s-backbone-renset50']
    # 需要绘制的指标
    lst_metrics = ['precision', 'recall', 'mAP@0.5', 'mAP50-95']
    # 各个指标同csv文件之间的索引对应关系
    dict_metrics_v5v8 = {'train/box_loss': 1, 'train/cls_loss': 2, 'train/dfl_loss': 3, 'precision': 4,
             'recall': 5, 'mAP@0.5': 6, 'mAP50-95': 7, 'val/box_loss': 8,
             'val/cls_loss': 9, 'val/dfl_loss': 10, 'lr/pg0': 11, 'lr/pg1': 12, 'lr/pg2': 13}

    # 绘制
    for metrics in lst_metrics:
        for modelname in lst_results_csv:
            res_path = os.path.join(root_path, modelname, 'results.csv')
            ext = res_path.split('.')[-1]
            # 适用yolov5、yolov8
            data = pd.read_csv(res_path, usecols=[dict_metrics_v5v8[metrics]]).values.ravel()
            x = range(len(data))
            plt.plot(x, data, label=modelname, linewidth='2')   # 线条粗细设为1

        # 添加x轴和y轴标签
        plt.xlabel('Epochs')
        plt.ylabel(metrics)
        plt.legend()
        plt.grid()
        # 显示图像
        plt.savefig(os.path.join(output_path, metrics + '.png'), dpi=900)   # dpi可设为300/600/900，表示存为更高清的矢量图
        plt.show()