import argparse import os import matplotlib.pyplot as plt import numpy as np from matplotlib.ticker import MultipleLocator from mmcv.ops import nms from mmdet.evaluation import bbox_overlaps from mmdet.utils import replace_cfg_vals, update_data_root from mmengine import Config, DictAction from mmengine.fileio import load from mmengine.registry import init_default_scope from mmengine.utils import ProgressBar from mmyolo.registry import DATASETS def parse_args(): parser = argparse.ArgumentParser( description='Generate confusion matrix from detection results') parser.add_argument('config', help='test config file path') parser.add_argument( 'prediction_path', help='prediction path where test .pkl result') parser.add_argument( 'save_dir', help='directory where confusion matrix will be saved') parser.add_argument( '--show', action='store_true', help='show confusion matrix') parser.add_argument( '--color-theme', default='plasma', help='theme of the matrix color map') parser.add_argument( '--score-thr', type=float, default=0.3, help='score threshold to filter detection bboxes') parser.add_argument( '--tp-iou-thr', type=float, default=0.5, help='IoU threshold to be considered as matched') parser.add_argument( '--nms-iou-thr', type=float, default=None, help='nms IoU threshold, only applied when users want to change the' 'nms IoU threshold.') parser.add_argument( '--cfg-options', nargs='+', action=DictAction, help='override some settings in the used config, the key-value pair ' 'in xxx=yyy format will be merged into config file. If the value to ' 'be overwritten is a list, it should be like key="[a,b]" or key=a,b ' 'It also allows nested list/tuple values, e.g. key="[(a,b),(c,d)]" ' 'Note that the quotation marks are necessary and that no white space ' 'is allowed.') args = parser.parse_args() return args def calculate_confusion_matrix(dataset, results, score_thr=0, nms_iou_thr=None, tp_iou_thr=0.5): """Calculate the confusion matrix. Args: dataset (Dataset): Test or val dataset. results (list[ndarray]): A list of detection results in each image. score_thr (float|optional): Score threshold to filter bboxes. Default: 0. nms_iou_thr (float|optional): nms IoU threshold, the detection results have done nms in the detector, only applied when users want to change the nms IoU threshold. Default: None. tp_iou_thr (float|optional): IoU threshold to be considered as matched. Default: 0.5. """ num_classes = len(dataset.metainfo['classes']) confusion_matrix = np.zeros(shape=[num_classes + 1, num_classes + 1]) assert len(dataset) == len(results) prog_bar = ProgressBar(len(results)) for idx, per_img_res in enumerate(results): res_bboxes = per_img_res['pred_instances'] gts = dataset.get_data_info(idx)['instances'] analyze_per_img_dets(confusion_matrix, gts, res_bboxes, score_thr, tp_iou_thr, nms_iou_thr) prog_bar.update() return confusion_matrix def analyze_per_img_dets(confusion_matrix, gts, result, score_thr=0, tp_iou_thr=0.5, nms_iou_thr=None): """Analyze detection results on each image. Args: confusion_matrix (ndarray): The confusion matrix, has shape (num_classes + 1, num_classes + 1). gt_bboxes (ndarray): Ground truth bboxes, has shape (num_gt, 4). gt_labels (ndarray): Ground truth labels, has shape (num_gt). result (ndarray): Detection results, has shape (num_classes, num_bboxes, 5). score_thr (float): Score threshold to filter bboxes. Default: 0. tp_iou_thr (float): IoU threshold to be considered as matched. Default: 0.5. nms_iou_thr (float|optional): nms IoU threshold, the detection results have done nms in the detector, only applied when users want to change the nms IoU threshold. Default: None. """ true_positives = np.zeros(len(gts)) gt_bboxes = [] gt_labels = [] for gt in gts: gt_bboxes.append(gt['bbox']) gt_labels.append(gt['bbox_label']) gt_bboxes = np.array(gt_bboxes) gt_labels = np.array(gt_labels) unique_label = np.unique(result['labels'].numpy()) for det_label in unique_label: mask = (result['labels'] == det_label) det_bboxes = result['bboxes'][mask].numpy() det_scores = result['scores'][mask].numpy() if nms_iou_thr: det_bboxes, _ = nms( det_bboxes, det_scores, nms_iou_thr, score_threshold=score_thr) ious = bbox_overlaps(det_bboxes[:, :4], gt_bboxes) for i, score in enumerate(det_scores): det_match = 0 if score >= score_thr: for j, gt_label in enumerate(gt_labels): if ious[i, j] >= tp_iou_thr: det_match += 1 if gt_label == det_label: true_positives[j] += 1 # TP confusion_matrix[gt_label, det_label] += 1 if det_match == 0: # BG FP confusion_matrix[-1, det_label] += 1 for num_tp, gt_label in zip(true_positives, gt_labels): if num_tp == 0: # FN confusion_matrix[gt_label, -1] += 1 def plot_confusion_matrix(confusion_matrix, labels, save_dir=None, show=True, title='Normalized Confusion Matrix', color_theme='plasma'): """Draw confusion matrix with matplotlib. Args: confusion_matrix (ndarray): The confusion matrix. labels (list[str]): List of class names. save_dir (str|optional): If set, save the confusion matrix plot to the given path. Default: None. show (bool): Whether to show the plot. Default: True. title (str): Title of the plot. Default: `Normalized Confusion Matrix`. color_theme (str): Theme of the matrix color map. Default: `plasma`. """ # normalize the confusion matrix per_label_sums = confusion_matrix.sum(axis=1)[:, np.newaxis] confusion_matrix = \ confusion_matrix.astype(np.float32) / per_label_sums * 100 num_classes = len(labels) fig, ax = plt.subplots( figsize=(0.5 * num_classes, 0.5 * num_classes * 0.8), dpi=180) cmap = plt.get_cmap(color_theme) im = ax.imshow(confusion_matrix, cmap=cmap) plt.colorbar(mappable=im, ax=ax) title_font = {'weight': 'bold', 'size': 12} ax.set_title(title, fontdict=title_font) label_font = {'size': 10} plt.ylabel('Ground Truth Label', fontdict=label_font) plt.xlabel('Prediction Label', fontdict=label_font) # draw locator xmajor_locator = MultipleLocator(1) xminor_locator = MultipleLocator(0.5) ax.xaxis.set_major_locator(xmajor_locator) ax.xaxis.set_minor_locator(xminor_locator) ymajor_locator = MultipleLocator(1) yminor_locator = MultipleLocator(0.5) ax.yaxis.set_major_locator(ymajor_locator) ax.yaxis.set_minor_locator(yminor_locator) # draw grid ax.grid(True, which='minor', linestyle='-') # draw label ax.set_xticks(np.arange(num_classes)) ax.set_yticks(np.arange(num_classes)) ax.set_xticklabels(labels) ax.set_yticklabels(labels) ax.tick_params( axis='x', bottom=False, top=True, labelbottom=False, labeltop=True) plt.setp( ax.get_xticklabels(), rotation=45, ha='left', rotation_mode='anchor') # draw confution matrix value for i in range(num_classes): for j in range(num_classes): ax.text( j, i, '{}%'.format( int(confusion_matrix[ i, j]) if not np.isnan(confusion_matrix[i, j]) else -1), ha='center', va='center', color='w', size=7) ax.set_ylim(len(confusion_matrix) - 0.5, -0.5) # matplotlib>3.1.1 fig.tight_layout() if save_dir is not None: plt.savefig( os.path.join(save_dir, 'confusion_matrix.png'), format='png') if show: plt.show() def main(): args = parse_args() cfg = Config.fromfile(args.config) # replace the ${key} with the value of cfg.key cfg = replace_cfg_vals(cfg) # update data root according to MMYOLO_DATASETS update_data_root(cfg) if args.cfg_options is not None: cfg.merge_from_dict(args.cfg_options) init_default_scope(cfg.get('default_scope', 'mmyolo')) results = load(args.prediction_path) if not os.path.exists(args.save_dir): os.makedirs(args.save_dir) dataset = DATASETS.build(cfg.test_dataloader.dataset) confusion_matrix = calculate_confusion_matrix(dataset, results, args.score_thr, args.nms_iou_thr, args.tp_iou_thr) plot_confusion_matrix( confusion_matrix, dataset.metainfo['classes'] + ('background', ), save_dir=args.save_dir, show=args.show, color_theme=args.color_theme) if __name__ == '__main__': main()