|
- # Copyright (c) OpenMMLab. All rights reserved.
- import argparse
- import os
-
- import matplotlib.pyplot as plt
- import mmcv
- import numpy as np
- from matplotlib.ticker import MultipleLocator
- from mmcv import Config, DictAction
-
- from mmseg.datasets import build_dataset
-
-
- def parse_args():
- parser = argparse.ArgumentParser(
- description='Generate confusion matrix from segmentation 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='winter',
- help='theme of the matrix color map')
- parser.add_argument(
- '--title',
- default='Normalized Confusion Matrix',
- help='title of the matrix color map')
- 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):
- """Calculate the confusion matrix.
-
- Args:
- dataset (Dataset): Test or val dataset.
- results (list[ndarray]): A list of segmentation results in each image.
- """
- n = len(dataset.CLASSES)
- confusion_matrix = np.zeros(shape=[n, n])
- assert len(dataset) == len(results)
- ignore_index = dataset.ignore_index
- prog_bar = mmcv.ProgressBar(len(results))
- for idx, per_img_res in enumerate(results):
- res_segm = per_img_res
- gt_segm = dataset.get_gt_seg_map_by_idx(idx).astype(int)
- gt_segm, res_segm = gt_segm.flatten(), res_segm.flatten()
- to_ignore = gt_segm == ignore_index
- gt_segm, res_segm = gt_segm[~to_ignore], res_segm[~to_ignore]
- inds = n * gt_segm + res_segm
- mat = np.bincount(inds, minlength=n**2).reshape(n, n)
- confusion_matrix += mat
- prog_bar.update()
- return confusion_matrix
-
-
- def plot_confusion_matrix(confusion_matrix,
- labels,
- save_dir=None,
- show=True,
- title='Normalized Confusion Matrix',
- color_theme='winter'):
- """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: `winter`.
- """
- # 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=(2 * num_classes, 2 * 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 confusion matrix value
- for i in range(num_classes):
- for j in range(num_classes):
- ax.text(
- j,
- i,
- '{}%'.format(
- round(confusion_matrix[i, j], 2
- ) 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], np.ndarray):
- pass
- 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)
- plot_confusion_matrix(
- confusion_matrix,
- dataset.CLASSES,
- save_dir=args.save_dir,
- show=args.show,
- title=args.title,
- color_theme=args.color_theme)
-
-
- if __name__ == '__main__':
- main()
|