Add confusion matrix tools analysis (#93)

tools/analysis_tools/confusion_matrix.py

@@ -0,0 +1,268 @@
+# Copyright (c) OpenMMLab. All rights reserved.
+import argparse
+import os
+
+import matplotlib.pyplot as plt
+import mmcv
+import numpy as np
+import torch
+from matplotlib.ticker import MultipleLocator
+from mmcv import Config, DictAction
+from mmcv.ops import nms_rotated
+from mmdet.datasets import build_dataset
+
+from mmrotate.core.bbox import rbbox_overlaps
+
+
+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.CLASSES)
+    confusion_matrix = np.zeros(shape=[num_classes + 1, num_classes + 1])
+    assert len(dataset) == len(results)
+    prog_bar = mmcv.ProgressBar(len(results))
+    for idx, per_img_res in enumerate(results):
+        if isinstance(per_img_res, tuple):
+            res_bboxes, _ = per_img_res
+        else:
+            res_bboxes = per_img_res
+        ann = dataset.get_ann_info(idx)
+        gt_bboxes = ann['bboxes']
+        labels = ann['labels']
+        analyze_per_img_dets(confusion_matrix, gt_bboxes, labels, res_bboxes,
+                             score_thr, tp_iou_thr, nms_iou_thr)
+        prog_bar.update()
+    return confusion_matrix
+
+
+def analyze_per_img_dets(confusion_matrix,
+                         gt_bboxes,
+                         gt_labels,
+                         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_like(gt_labels)
+    for det_label, det_bboxes in enumerate(result):
+        det_bboxes = torch.from_numpy(det_bboxes).float()
+        gt_bboxes = torch.from_numpy(gt_bboxes).float()
+        if nms_iou_thr:
+            det_bboxes, _ = nms_rotated(
+                det_bboxes[:, :5],
+                det_bboxes[:, -1],
+                nms_iou_thr,
+                score_threshold=score_thr)
+        ious = rbbox_overlaps(det_bboxes[:, :5], gt_bboxes)
+        for i, det_bbox in enumerate(det_bboxes):
+            score = det_bbox[5]
+            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)
+    if args.cfg_options is not None:
+        cfg.merge_from_dict(args.cfg_options)
+
+    results = mmcv.load(args.prediction_path)
+    assert isinstance(results, list)
+    if isinstance(results[0], list):
+        pass
+    elif isinstance(results[0], tuple):
+        results = [result[0] for result in results]
+    else:
+        raise TypeError('invalid type of prediction results')
+
+    if isinstance(cfg.data.test, dict):
+        cfg.data.test.test_mode = True
+    elif isinstance(cfg.data.test, list):
+        for ds_cfg in cfg.data.test:
+            ds_cfg.test_mode = True
+    dataset = build_dataset(cfg.data.test)
+
+    confusion_matrix = calculate_confusion_matrix(dataset, results,
+                                                  args.score_thr,
+                                                  args.nms_iou_thr,
+                                                  args.tp_iou_thr)
+    plot_confusion_matrix(
+        confusion_matrix,
+        dataset.CLASSES + ('background', ),
+        save_dir=args.save_dir,
+        show=args.show)
+
+
+if __name__ == '__main__':
+    main()