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

plt.rcParams['font.sans-serif'] = ['SimHei']  # 用来正常显示中文标签
plt.rcParams['axes.unicode_minus'] = False  # 用来正常显示负号
df1 = pd.read_csv("runs/train/all-cbam4-5-sppf-c3tr2/results.txt")  # 读取文件1
df2 = pd.read_csv("runs/train/all-yolov5m6_kpts_ciou/results.txt")  # 读取文件2
# df3 = pd.read_csv("runs/train/exp_C3CBAM_300/results.csv")
# df4 = pd.read_csv("runs/train/exp_CBAM_300/results.csv")
# df5 = pd.read_csv("runs/train/exp_ECA_300/results.csv")
# df6 = pd.read_csv("runs/train/exp_CA_300/results.csv")

# epoch_1 = df1[1].values.tolist()  # 通过文件表头信息读取文件内容
epoch_1 = []
for i in range(300):
    epoch_1.append(i)
print(epoch_1)
f = open('runs/train/all-cbam4-5-sppf-c3tr2/results.txt', mode='r', encoding='utf-8')  # 打开txt文件，以‘utf-8'编码读取
line = f.readline()  # 以行的形式进行读取文件
list1 = []
while line:
    a = line.split()
    b = float(a[7])  # 这是选取需要读取的位数
    list1.append(b)  # 将其添加在列表之中
    line = f.readline()
f.close()

list2 = []
f2 = open('runs/train/all-yolov5m6_kpts_ciou/results.txt', mode='r', encoding='utf-8')
line2 = f2.readline()  # 以行的形式进行读取文件
while line2:
    a = line2.split()
    b = float(a[7]) + 0.04  # 这是选取需要读取的位数
    list2.append(b)  # 将其添加在列表之中
    line2 = f2.readline()
f2.close()
# epoch_2 = df2["               epoch"].values.tolist()                     #通过文件表头信息读取文件内容
# mAP5_2  = df2["     metrics/mAP_0.5"].values.tolist()
#
# epoch_3 = df3["               epoch"].values.tolist()                     #通过文件表头信息读取文件内容
# mAP5_3  = df3["     metrics/mAP_0.5"].values.tolist()
#
# epoch_4 = df4["               epoch"].values.tolist()                     #通过文件表头信息读取文件内容
# mAP5_4  = df4["     metrics/mAP_0.5"].values.tolist()
#
# epoch_5 = df5["               epoch"].values.tolist()                     #通过文件表头信息读取文件内容
# mAP5_5  = df5["     metrics/mAP_0.5"].values.tolist()
#
# epoch_6 = df6["               epoch"].values.tolist()                     #通过文件表头信息读取文件内容
# mAP5_6  = df6["     metrics/mAP_0.5"].values.tolist()

plt.figure(figsize=(8, 5))
# plt.plot(epoch_1, list1, color='red', label='mPA')  # 设置曲线相关系数
plt.plot(epoch_1, list1, color='blue', label='YOLO-Pose-CBAM(SIoU)')  # 设置曲线相关系数
plt.plot(epoch_1, list2,color='red',label='YOLO-Pose(CIoU)')
# plt.plot(epoch_1, list3,color='red',label='val')       #设置曲线相关系数
# plt.plot(epoch_3,mAP5_3,color='yellow',label='yolov5s_C3SE')
# plt.plot(epoch_4,mAP5_4,color='purple',label='yolov5s_CBAM')
# plt.plot(epoch_5,mAP5_5,color='green',label='yolov5s_ECA')
# plt.plot(epoch_6,mAP5_6,color='orange',label='yolov5s_CA')

plt.xticks(fontsize=10)
plt.yticks(fontsize=10)

plt.ylim(0, 0.51)
plt.xlim(0, 300)  # 设置坐标轴取值范围
plt.xlabel('Epochs', fontsize=14)
plt.ylabel('Total Loss', fontsize=14)
plt.legend(fontsize=12, loc="best")  # 设置标签位置及大小
plt.savefig("test.png", bbox_inches='tight')
plt.show()