训练结束后分析训练数据集
- 在训练完一个分类模型后,为了分析训练模型在验证/训练数据集上不同类别效果,往往要计算其不同类别的精确率和召回率,以及检测错误的图片有哪些,错误图片被误检为了什么。
- 首先文件夹格式为
- 直接贴出代码,实现了对不同类别检测结果并保存错误结果,结果如下显示
![在这里插入图片描述](https://i-blog.csdnimg.cn/direct/06232276a5fb49dd8172a2e9bec3315a.png)
from ultralytics import YOLO
import matplotlib
import os
from sklearn.metrics import accuracy_score, precision_score,f1_score,recall_score
from sklearn.metrics import classification_report
import cv2
model = YOLO('runs/classify/train20/weights/best.pt')
names = {0: 'CCC', 1: 'DDD', 2: 'Other', 3: 'QQQ'}
base_path = '/home/train'
dict_names = {v: k for k, v in names.items()}
real_labels = []
pre_labels = []
for i in os.listdir(base_path):
label = dict_names[i]
for j in os.listdir(os.path.join(base_path, i)):
img_path = os.path.join(base_path, i, j)
filename = img_path.split('/')[-1]
res = model.predict(img_path)[0]
if res.probs.top1 != label:
image = cv2.imread(img_path)
name = i
resname = names[res.probs.top1]
print("预测错误:", filename,"label:",name,"res:",resname)
output_path = os.path.join("train", name, resname, filename)
os.makedirs(os.path.dirname(output_path), exist_ok=True)
cv2.imwrite(output_path,image)
real_labels.append(label)
pre_labels.append(res.probs.top1)
print("每个类别的精确率、召回率和F1-Score:")
print(classification_report(real_labels, pre_labels, target_names=list(names.values())))
参考:
1
2