|
|
|
|
|
import argparse
|
|
|
import logging
|
|
|
import tempfile
|
|
|
|
|
|
import torch
|
|
|
from mmengine import dump, list_from_file, load
|
|
|
from mmengine.config import Config, DictAction
|
|
|
from mmengine.evaluator import Evaluator
|
|
|
from mmengine.runner import Runner
|
|
|
|
|
|
from mmaction.evaluation import ConfusionMatrix
|
|
|
from mmaction.registry import DATASETS
|
|
|
from mmaction.utils import register_all_modules
|
|
|
|
|
|
|
|
|
def parse_args():
|
|
|
parser = argparse.ArgumentParser(
|
|
|
description='Eval a checkpoint and draw the confusion matrix.')
|
|
|
parser.add_argument('config', help='test config file path')
|
|
|
parser.add_argument(
|
|
|
'ckpt_or_result',
|
|
|
type=str,
|
|
|
help='The checkpoint file (.pth) or '
|
|
|
'dumpped predictions pickle file (.pkl).')
|
|
|
parser.add_argument('--out', help='the file to save the confusion matrix.')
|
|
|
parser.add_argument(
|
|
|
'--show',
|
|
|
action='store_true',
|
|
|
help='whether to display the metric result by matplotlib if supports.')
|
|
|
parser.add_argument(
|
|
|
'--show-path', type=str, help='Path to save the visualization image.')
|
|
|
parser.add_argument(
|
|
|
'--include-values',
|
|
|
action='store_true',
|
|
|
help='To draw the values in the figure.')
|
|
|
parser.add_argument('--label-file', default=None, help='Labelmap file')
|
|
|
parser.add_argument(
|
|
|
'--target-classes',
|
|
|
type=int,
|
|
|
nargs='+',
|
|
|
default=[],
|
|
|
help='Selected classes to evaluate, and remains will be neglected')
|
|
|
parser.add_argument(
|
|
|
'--cmap',
|
|
|
type=str,
|
|
|
default='viridis',
|
|
|
help='The color map to use. Defaults to "viridis".')
|
|
|
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 main():
|
|
|
args = parse_args()
|
|
|
|
|
|
|
|
|
|
|
|
register_all_modules(init_default_scope=False)
|
|
|
|
|
|
|
|
|
cfg = Config.fromfile(args.config)
|
|
|
if args.cfg_options is not None:
|
|
|
cfg.merge_from_dict(args.cfg_options)
|
|
|
|
|
|
if args.ckpt_or_result.endswith('.pth'):
|
|
|
|
|
|
cfg.test_evaluator = dict(type='ConfusionMatrix')
|
|
|
|
|
|
cfg.load_from = str(args.ckpt_or_result)
|
|
|
|
|
|
with tempfile.TemporaryDirectory() as tmpdir:
|
|
|
cfg.work_dir = tmpdir
|
|
|
runner = Runner.from_cfg(cfg)
|
|
|
classes = runner.test_loop.dataloader.dataset.metainfo.get(
|
|
|
'classes')
|
|
|
cm = runner.test()['confusion_matrix/result']
|
|
|
logging.shutdown()
|
|
|
else:
|
|
|
predictions = load(args.ckpt_or_result)
|
|
|
evaluator = Evaluator(ConfusionMatrix())
|
|
|
metrics = evaluator.offline_evaluate(predictions, None)
|
|
|
cm = metrics['confusion_matrix/result']
|
|
|
try:
|
|
|
|
|
|
dataset = DATASETS.build({
|
|
|
**cfg.test_dataloader.dataset, 'pipeline': []
|
|
|
})
|
|
|
classes = dataset.metainfo.get('classes')
|
|
|
except Exception:
|
|
|
classes = None
|
|
|
|
|
|
if args.label_file is not None:
|
|
|
classes = list_from_file(args.label_file)
|
|
|
if classes is None:
|
|
|
num_classes = cm.shape[0]
|
|
|
classes = list(range(num_classes))
|
|
|
|
|
|
if args.target_classes:
|
|
|
assert len(args.target_classes) > 1, \
|
|
|
'please ensure select more than one class'
|
|
|
target_idx = torch.tensor(args.target_classes)
|
|
|
cm = cm[target_idx][:, target_idx]
|
|
|
classes = [classes[idx] for idx in target_idx]
|
|
|
|
|
|
if args.out is not None:
|
|
|
dump(cm, args.out)
|
|
|
|
|
|
if args.show or args.show_path is not None:
|
|
|
fig = ConfusionMatrix.plot(
|
|
|
cm,
|
|
|
show=args.show,
|
|
|
classes=classes,
|
|
|
include_values=args.include_values,
|
|
|
cmap=args.cmap)
|
|
|
if args.show_path is not None:
|
|
|
fig.savefig(args.show_path)
|
|
|
print(f'The confusion matrix is saved at {args.show_path}.')
|
|
|
|
|
|
|
|
|
if __name__ == '__main__':
|
|
|
main()
|
|
|
|
