mirror of
https://github.com/open-mmlab/mmyolo.git
synced 2025-06-03 15:00:20 +08:00
274 lines
9.7 KiB
Python
274 lines
9.7 KiB
Python
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()
|