yolov5混淆矩阵图如何绘制

### 绘制YOLOv5的混淆矩阵图 为了在YOLOv5中绘制混淆矩阵图，通常需要利用`val.py`脚本中的功能并扩展其代码以保存预测结果和真实标签。之后可以使用这些数据来构建混淆矩阵，并借助matplotlib库进行可视化。 #### 修改 `val.py` 首先，在`val.py`中找到合适的位置加入代码片段以便收集每一轮验证过程中产生的预测框与实际标注之间的匹配情况。这一步骤涉及修改或增加函数调用来记录每次推理的结果： ```python from pathlib import Path import numpy as np import torch from utils.metrics import fitness from utils.general import xywh2xyxy, box_iou from collections import defaultdict def process_batch(detections, labels, iouv): """ Process a batch of detections and ground truth labels. Args: detections (Tensor[N, 6]): Detected boxes with confidence scores and class indices. labels (Tensor[M, 5]): Ground truth boxes including classes. iouv (Tensor[J]): IoU vector for evaluation at different thresholds. Returns: dict: Dictionary containing statistics about TP/FP/TN/FN per class. """ stats = {'tp': [], 'fp': []} correct = [] detected = [] if len(labels) == 0 or len(detections) == 0: return {k: v.tolist() for k, v in stats.items()}, correct target_classes_o = labels[:, 0].tolist() pred_boxes = detections[:, :4] pred_scores = detections[:, 4] pred_labels = detections[:, 5] target_boxes = labels[:, 1:] target_labels = labels[:, 0] # Convert to absolute coordinates scale_coords((h, w), pred_boxes, img_shape=(img_size, img_size)).round() matches = defaultdict(list) for di, d in enumerate(pred_boxes): ci = int(pred_labels[di]) jidx = (pred_labels[di] == target_labels).nonzero(as_tuple=False) if jidx.size(0) == 0: continue ious = box_iou(d.unsqueeze(0), target_boxes[jidx]).squeeze() best_i = ious.argmax() if ious.numel() > 0 else None if best_i is not None and ious[best_i] >= iouv[-1]: det_idx = jidx[best_i][0] if det_idx.item() not in detected: detected.append(det_idx.item()) correct.append([ci, True]) # 正确检测到目标 matches[target_classes_o[int(det_idx)]].append(di) elif pred_scores[di] > pred_scores[matches[target_classes_o[int(det_idx)]][-1]]: # 更高得分替换之前的假阳性为真阳性 old_di = matches[target_classes_o[int(det_idx)]][-1] correct[old_di] = [target_classes_o[int(det_idx)], False] correct.append([ci, True]) matches[target_classes_o[int(det_idx)]][-1] = di else: correct.append([ci, False]) # 假阳性的其他情况 else: correct.append([ci, False]) # 完全错误的情况 for cls_id, match_ids in matches.items(): unmatched_targets = set(range(len(target_classes_o))) - set(match_ids) for ut in unmatched_targets: correct.append([cls_id, False]) # 对应类别的未命中（即漏检） tcls = list(map(int, target_classes_o)) pcls = [int(x) for x in pred_labels.cpu().numpy()] tp = sum(c[1] for c in correct if c[0] in tcls) fp = sum(not c[1] for c in correct if c[0] in pcls) stats['tp'].extend([tp]*len(tcls)) stats['fp'].extend([fp]*len(pcls)) return {k: v.tolist() for k, v in stats.items()}, correct ``` 上述代码实现了对一批次内所有预测框与真实框之间关系的处理逻辑[^3]。接下来还需要调整主循环部分以累积统计信息直到整个验证集遍历完毕。 #### 构建混淆矩阵 完成以上准备工作后，可以在程序结束前添加如下Python代码段用于创建混淆矩阵图表： ```python import seaborn as sns import matplotlib.pyplot as plt from sklearn.metrics import confusion_matrix import pandas as pd # Assuming you have collected all predictions and true labels during validation phase into two lists called preds and trues respectively cm = confusion_matrix(trues, preds) df_cm = pd.DataFrame(cm, index=[i for i in "ABCDEFGHIJKL"], columns=[i for i in "ABCDEFGHIJKL"]) plt.figure(figsize=(10,7)) sns.heatmap(df_cm, annot=True) plt.title('Confusion Matrix') plt.ylabel('Actual Label') plt.xlabel('Predicted Label') save_path = str(Path(save_dir) / f'confusion_matrix.png') # Adjust according to your directory structure plt.savefig(save_path) print(f'Saved confusion matrix figure to "{save_path}"') ``` 这段代码假设已经有一个名为`trues`的真实标签列表和另一个名为`preds`的预测标签列表被填充好了。通过这两个列表可以直接调用scikit-learn提供的`confusion_matrix()`方法得到混淆矩阵数组，再配合seaborn库画出热力图形式展示出来[^2]。