【YOLO系列绘图】根据v5、v8、v7训练后生成的result文件用matplotlib进行绘图

效果图

代码

• 需要更改的地方：仅为result_dict中的内容，每行表示的是模型名称和该模型训练得到的result文件地址
• v5、v8都是csv格式的，v7是txt格式的

import matplotlib.pyplot as plt
import pandas as pd
import numpy as np

if __name__ == '__main__':
    # 列出待获取数据内容的文件位置
    # v5、v8都是csv格式的，v7是txt格式的
    result_dict = {
    
    
        'YOLOv5m': r'C:\Users\dadandan\Desktop\workspace_temp\NeimengA_train_resultcsv\yolov5m\results.csv',
        'YOLOv7': r'C:\Users\dadandan\Desktop\workspace_temp\NeimengA_train_resultcsv\yolov7\results.txt',
        'YOLOv7-tiny': r'C:\Users\dadandan\Desktop\workspace_temp\NeimengA_train_resultcsv\yolov7-tiny\results.txt',
        'YOLOv7-tiny-large': r'C:\Users\dadandan\Desktop\workspace_temp\NeimengA_train_resultcsv\yolov7-tiny-large\results.txt',
        'YOLOv7-tiny-PConv': r'C:\Users\dadandan\Desktop\workspace_temp\NeimengA_train_resultcsv\yolov7-tiny-PConv\results.txt',
        'YOLOv7-tiny-GhostNetv2': r'C:\Users\dadandan\Desktop\workspace_temp\NeimengA_train_resultcsv\yolov7-tiny-GhostNetv2\results.txt',
        'YOLOv8s': r'C:\Users\dadandan\Desktop\workspace_temp\NeimengA_train_resultcsv\yolov8s\results.csv'
    }

    # 绘制map50
    for modelname in result_dict:
        res_path = result_dict[modelname]
        ext = res_path.split('.')[-1]
        if ext == 'csv':
            data = pd.read_csv(res_path, usecols=[6]).values.ravel()
        else:
            with open(res_path, 'r') as f:
                datalist = f.readlines()
                data = []
                for d in datalist:
                    data.append(float(d.strip().split()[10]))
                data = np.array(data)
        x = range(len(data))
        plt.plot(x, data, label=modelname, linewidth='2')

    # 添加x轴和y轴标签
    plt.xlabel('Epochs')
    plt.ylabel('[email protected]')
    plt.legend()
    plt.grid()
    # 保存并显示图像
    plt.savefig("mAP50.png")
    plt.show()


    # 绘制map50-95
    for modelname in result_dict:
        res_path = result_dict[modelname]
        ext = res_path.split('.')[-1]
        if ext == 'csv':
            data = pd.read_csv(res_path, usecols=[7]).values.ravel()
        else:
            with open(res_path, 'r') as f:
                datalist = f.readlines()
                data = []
                for d in datalist:
                    data.append(float(d.strip().split()[11]))
                data = np.array(data)
        x = range(len(data))
        plt.plot(x, data, label=modelname, linewidth='2')

    # 添加x轴和y轴标签
    plt.xlabel('Epochs')
    plt.ylabel('[email protected]:0.95')
    plt.legend()
    plt.grid()
    # 保存并显示图像
    plt.savefig("mAP50-95.png")
    plt.show()