diff --git a/configs/hmr/resnet50_hmr_pw3d.py b/configs/hmr/resnet50_hmr_pw3d.py index 1fc2ff83..83a2b60a 100644 --- a/configs/hmr/resnet50_hmr_pw3d.py +++ b/configs/hmr/resnet50_hmr_pw3d.py @@ -1,6 +1,8 @@ _base_ = ['../_base_/default_runtime.py'] use_adversarial_train = True +# evaluate +evaluation = dict(metric=['pa-mpjpe', 'mpjpe']) # optimizer optimizer = dict( backbone=dict(type='Adam', lr=2.5e-4), diff --git a/configs/hybrik/resnet34_hybrik_mixed.py b/configs/hybrik/resnet34_hybrik_mixed.py index ab68bae4..ebefe2ec 100644 --- a/configs/hybrik/resnet34_hybrik_mixed.py +++ b/configs/hybrik/resnet34_hybrik_mixed.py @@ -1,5 +1,7 @@ _base_ = ['../_base_/default_runtime.py'] +# evaluate +evaluation = dict(metric=['pa-mpjpe', 'mpjpe']) # optimizer optimizer = dict(type='Adam', lr=1e-3, weight_decay=0) optimizer_config = dict(grad_clip=None) @@ -11,7 +13,7 @@ interval=50, hooks=[ dict(type='TextLoggerHook'), - # dict(type='TensorboardLoggerHook') + # dict(type='TensorboardLoggerHook') ]) img_res = 256 @@ -166,7 +168,16 @@ partition=[0.4, 0.1, 0.5]), test=dict( type=dataset_type, + body_model=dict( + type='GenderedSMPL', model_path='data/body_models/smpl'), dataset_name='pw3d', data_prefix='data', pipeline=test_pipeline, - ann_file='hybrik_pw3d_test.npz')) + ann_file='hybrik_pw3d_test.npz'), + val=dict( + type=dataset_type, + dataset_name='pw3d', + data_prefix='data', + pipeline=test_pipeline, + ann_file='hybrik_pw3d_test.npz'), +) diff --git a/configs/spin/resnet50_spin_pw3d.py b/configs/spin/resnet50_spin_pw3d.py index 68ed638c..d5ed1851 100644 --- a/configs/spin/resnet50_spin_pw3d.py +++ b/configs/spin/resnet50_spin_pw3d.py @@ -1,6 +1,9 @@ _base_ = ['../_base_/default_runtime.py'] use_adversarial_train = True +# evaluate +evaluation = dict(metric=['pa-mpjpe', 'mpjpe']) + img_res = 224 body_model = dict( diff --git a/configs/vibe/resnet50_vibe_pw3d.py b/configs/vibe/resnet50_vibe_pw3d.py index 9a7c8f16..d6ea5290 100644 --- a/configs/vibe/resnet50_vibe_pw3d.py +++ b/configs/vibe/resnet50_vibe_pw3d.py @@ -1,6 +1,9 @@ _base_ = ['../_base_/default_runtime.py'] use_adversarial_train = True +# evaluate +evaluation = dict(metric=['pa-mpjpe', 'mpjpe']) + # optimizer optimizer = dict( neck=dict(type='Adam', lr=2.5e-4), head=dict(type='Adam', lr=2.5e-4)) diff --git a/docs/getting_started.md b/docs/getting_started.md index 3c64c0c8..015da079 100644 --- a/docs/getting_started.md +++ b/docs/getting_started.md @@ -110,11 +110,11 @@ We provide pretrained models in the respective method folders in [config](https: ### Evaluate with a single GPU / multiple GPUs ```shell -python tools/test.py ${CONFIG} --work-dir=${WORK_DIR} ${CHECKPOINT} +python tools/test.py ${CONFIG} --work-dir=${WORK_DIR} ${CHECKPOINT} --metrics=${METRICS} ``` Example: ```shell -python tools/test.py configs/hmr/resnet50_hmr_pw3d.py --work-dir=work_dirs/hmr work_dirs/hmr/latest.pth +python tools/test.py configs/hmr/resnet50_hmr_pw3d.py --work-dir=work_dirs/hmr work_dirs/hmr/latest.pth --metrics pa-mpjpe mpjpe ``` ### Evaluate with slurm @@ -122,11 +122,11 @@ python tools/test.py configs/hmr/resnet50_hmr_pw3d.py --work-dir=work_dirs/hmr w If you can run MMHuman3D on a cluster managed with [slurm](https://slurm.schedmd.com/), you can use the script `slurm_test.sh`. ```shell -./tools/slurm_test.sh ${PARTITION} ${JOB_NAME} ${CONFIG} ${WORK_DIR} ${CHECKPOINT} +./tools/slurm_test.sh ${PARTITION} ${JOB_NAME} ${CONFIG} ${WORK_DIR} ${CHECKPOINT} --metrics ${METRICS} ``` Example: ```shell -./tools/slurm_test.sh my_partition test_hmr configs/hmr/resnet50_hmr_pw3d.py work_dirs/hmr work_dirs/hmr/latest.pth 8 +./tools/slurm_test.sh my_partition test_hmr configs/hmr/resnet50_hmr_pw3d.py work_dirs/hmr work_dirs/hmr/latest.pth 8 --metrics pa-mpjpe mpjpe ``` diff --git a/mmhuman3d/apis/test.py b/mmhuman3d/apis/test.py index c8736e6c..332c328b 100644 --- a/mmhuman3d/apis/test.py +++ b/mmhuman3d/apis/test.py @@ -5,18 +5,12 @@ import time import mmcv -import numpy as np import torch import torch.distributed as dist -from mmcv.image import tensor2imgs from mmcv.runner import get_dist_info -def single_gpu_test(model, - data_loader, - show=False, - out_dir=None, - **show_kwargs): +def single_gpu_test(model, data_loader): """Test with single gpu.""" model.eval() results = [] @@ -32,40 +26,6 @@ def single_gpu_test(model, else: results.append(result) - if show or out_dir: - scores = np.vstack(result) - pred_score = np.max(scores, axis=1) - pred_label = np.argmax(scores, axis=1) - pred_class = [model.CLASSES[lb] for lb in pred_label] - - img_metas = data['img_metas'].data[0] - imgs = tensor2imgs(data['img'], **img_metas[0]['img_norm_cfg']) - assert len(imgs) == len(img_metas) - - for i, (img, img_meta) in enumerate(zip(imgs, img_metas)): - h, w, _ = img_meta['img_shape'] - img_show = img[:h, :w, :] - - ori_h, ori_w = img_meta['ori_shape'][:-1] - img_show = mmcv.imresize(img_show, (ori_w, ori_h)) - - if out_dir: - out_file = osp.join(out_dir, img_meta['ori_filename']) - else: - out_file = None - - result_show = { - 'pred_score': pred_score[i], - 'pred_label': pred_label[i], - 'pred_class': pred_class[i] - } - model.module.show_result( - img_show, - result_show, - show=show, - out_file=out_file, - **show_kwargs) - if 'img' in data.keys(): batch_size = data['img'].size(0) else: diff --git a/mmhuman3d/apis/train.py b/mmhuman3d/apis/train.py index ad017cda..aca92d5a 100644 --- a/mmhuman3d/apis/train.py +++ b/mmhuman3d/apis/train.py @@ -10,9 +10,9 @@ OptimizerHook, build_runner, ) -from mmcv.runner.hooks import DistEvalHook, EvalHook from mmhuman3d.core.distributed_wrapper import DistributedDataParallelWrapper +from mmhuman3d.core.evaluation import DistEvalHook, EvalHook from mmhuman3d.core.optimizer import build_optimizers from mmhuman3d.data.datasets import build_dataloader, build_dataset from mmhuman3d.utils import get_root_logger @@ -156,7 +156,6 @@ def train_model(model, round_up=True) eval_cfg = cfg.get('evaluation', {}) eval_cfg['by_epoch'] = cfg.runner['type'] != 'IterBasedRunner' - eval_cfg['work_dir'] = cfg.work_dir eval_hook = DistEvalHook if distributed else EvalHook runner.register_hook(eval_hook(val_dataloader, **eval_cfg)) diff --git a/mmhuman3d/core/evaluation/__init__.py b/mmhuman3d/core/evaluation/__init__.py index 73fe2777..1923b8ad 100644 --- a/mmhuman3d/core/evaluation/__init__.py +++ b/mmhuman3d/core/evaluation/__init__.py @@ -1,7 +1,16 @@ -from mmhuman3d.core.evaluation import mesh_eval, mpjpe +from mmhuman3d.core.evaluation import mesh_eval +from mmhuman3d.core.evaluation.eval_hooks import DistEvalHook, EvalHook +from mmhuman3d.core.evaluation.eval_utils import ( + keypoint_3d_auc, + keypoint_3d_pck, + keypoint_accel_error, + keypoint_mpjpe, + vertice_pve, +) from mmhuman3d.core.evaluation.mesh_eval import compute_similarity_transform -from mmhuman3d.core.evaluation.mpjpe import keypoint_mpjpe __all__ = [ - 'compute_similarity_transform', 'keypoint_mpjpe', 'mesh_eval', 'mpjpe' + 'compute_similarity_transform', 'keypoint_mpjpe', 'mesh_eval', + 'DistEvalHook', 'EvalHook', 'vertice_pve', 'keypoint_3d_pck', + 'keypoint_3d_auc', 'keypoint_accel_error' ] diff --git a/mmhuman3d/core/evaluation/eval_hooks.py b/mmhuman3d/core/evaluation/eval_hooks.py new file mode 100644 index 00000000..fccc6eca --- /dev/null +++ b/mmhuman3d/core/evaluation/eval_hooks.py @@ -0,0 +1,139 @@ +# Copyright (c) OpenMMLab. All rights reserved. +import tempfile +import warnings + +from mmcv.runner import DistEvalHook as BaseDistEvalHook +from mmcv.runner import EvalHook as BaseEvalHook + +MMHUMAN3D_GREATER_KEYS = ['3dpck', 'pa-3dpck', '3dauc', 'pa-3dauc'] +MMHUMAN3D_LESS_KEYS = ['mpjpe', 'pa-mpjpe', 'pve'] + + +class EvalHook(BaseEvalHook): + + def __init__(self, + dataloader, + start=None, + interval=1, + by_epoch=True, + save_best=None, + rule=None, + test_fn=None, + greater_keys=MMHUMAN3D_GREATER_KEYS, + less_keys=MMHUMAN3D_LESS_KEYS, + **eval_kwargs): + if test_fn is None: + from mmhuman3d.apis import single_gpu_test + test_fn = single_gpu_test + + # remove "gpu_collect" from eval_kwargs + if 'gpu_collect' in eval_kwargs: + warnings.warn( + '"gpu_collect" will be deprecated in EvalHook.' + 'Please remove it from the config.', DeprecationWarning) + _ = eval_kwargs.pop('gpu_collect') + + # update "save_best" according to "key_indicator" and remove the + # latter from eval_kwargs + if 'key_indicator' in eval_kwargs or isinstance(save_best, bool): + warnings.warn( + '"key_indicator" will be deprecated in EvalHook.' + 'Please use "save_best" to specify the metric key,' + 'e.g., save_best="pa-mpjpe".', DeprecationWarning) + + key_indicator = eval_kwargs.pop('key_indicator', None) + if save_best is True and key_indicator is None: + raise ValueError('key_indicator should not be None, when ' + 'save_best is set to True.') + save_best = key_indicator + + super().__init__(dataloader, start, interval, by_epoch, save_best, + rule, test_fn, greater_keys, less_keys, **eval_kwargs) + + def evaluate(self, runner, results): + + with tempfile.TemporaryDirectory() as tmp_dir: + eval_res = self.dataloader.dataset.evaluate( + results, + res_folder=tmp_dir, + logger=runner.logger, + **self.eval_kwargs) + + for name, val in eval_res.items(): + runner.log_buffer.output[name] = val + runner.log_buffer.ready = True + + if self.save_best is not None: + if self.key_indicator == 'auto': + self._init_rule(self.rule, list(eval_res.keys())[0]) + + return eval_res[self.key_indicator] + + return None + + +class DistEvalHook(BaseDistEvalHook): + + def __init__(self, + dataloader, + start=None, + interval=1, + by_epoch=True, + save_best=None, + rule=None, + test_fn=None, + greater_keys=MMHUMAN3D_GREATER_KEYS, + less_keys=MMHUMAN3D_LESS_KEYS, + broadcast_bn_buffer=True, + tmpdir=None, + gpu_collect=False, + **eval_kwargs): + + if test_fn is None: + from mmhuman3d.apis import multi_gpu_test + test_fn = multi_gpu_test + + # update "save_best" according to "key_indicator" and remove the + # latter from eval_kwargs + if 'key_indicator' in eval_kwargs or isinstance(save_best, bool): + warnings.warn( + '"key_indicator" will be deprecated in EvalHook.' + 'Please use "save_best" to specify the metric key,' + 'e.g., save_best="pa-mpjpe".', DeprecationWarning) + + key_indicator = eval_kwargs.pop('key_indicator', None) + if save_best is True and key_indicator is None: + raise ValueError('key_indicator should not be None, when ' + 'save_best is set to True.') + save_best = key_indicator + + super().__init__(dataloader, start, interval, by_epoch, save_best, + rule, test_fn, greater_keys, less_keys, + broadcast_bn_buffer, tmpdir, gpu_collect, + **eval_kwargs) + + def evaluate(self, runner, results): + """Evaluate the results. + + Args: + runner (:obj:`mmcv.Runner`): The underlined training runner. + results (list): Output results. + """ + with tempfile.TemporaryDirectory() as tmp_dir: + eval_res = self.dataloader.dataset.evaluate( + results, + res_folder=tmp_dir, + logger=runner.logger, + **self.eval_kwargs) + + for name, val in eval_res.items(): + runner.log_buffer.output[name] = val + runner.log_buffer.ready = True + + if self.save_best is not None: + if self.key_indicator == 'auto': + # infer from eval_results + self._init_rule(self.rule, list(eval_res.keys())[0]) + return eval_res[self.key_indicator] + + return None diff --git a/mmhuman3d/core/evaluation/eval_utils.py b/mmhuman3d/core/evaluation/eval_utils.py new file mode 100644 index 00000000..e1d001b2 --- /dev/null +++ b/mmhuman3d/core/evaluation/eval_utils.py @@ -0,0 +1,199 @@ +# Copyright (c) OpenMMLab. All rights reserved. +import numpy as np + +from .mesh_eval import compute_similarity_transform + + +def keypoint_mpjpe(pred, gt, mask, alignment='none'): + """Calculate the mean per-joint position error (MPJPE) and the error after + rigid alignment with the ground truth (PA-MPJPE). + batch_size: N + num_keypoints: K + keypoint_dims: C + Args: + pred (np.ndarray[N, K, C]): Predicted keypoint location. + gt (np.ndarray[N, K, C]): Groundtruth keypoint location. + mask (np.ndarray[N, K]): Visibility of the target. False for invisible + joints, and True for visible. Invisible joints will be ignored for + accuracy calculation. + alignment (str, optional): method to align the prediction with the + groundtruth. Supported options are: + - ``'none'``: no alignment will be applied + - ``'scale'``: align in the least-square sense in scale + - ``'procrustes'``: align in the least-square sense in scale, + rotation and translation. + Returns: + tuple: A tuple containing joint position errors + - mpjpe (float|np.ndarray[N]): mean per-joint position error. + - pa-mpjpe (float|np.ndarray[N]): mpjpe after rigid alignment with the + ground truth + """ + assert mask.any() + + if alignment == 'none': + pass + elif alignment == 'procrustes': + pred = np.stack([ + compute_similarity_transform(pred_i, gt_i) + for pred_i, gt_i in zip(pred, gt) + ]) + elif alignment == 'scale': + pred_dot_pred = np.einsum('nkc,nkc->n', pred, pred) + pred_dot_gt = np.einsum('nkc,nkc->n', pred, gt) + scale_factor = pred_dot_gt / pred_dot_pred + pred = pred * scale_factor[:, None, None] + else: + raise ValueError(f'Invalid value for alignment: {alignment}') + + error = np.linalg.norm(pred - gt, ord=2, axis=-1)[mask].mean() + + return error + + +def keypoint_accel_error(gt, pred, mask=None): + """Computes acceleration error: + + Note that for each frame that is not visible, three entries in the + acceleration error should be zero'd out. + Args: + gt (Nx14x3). + pred (Nx14x3). + mask (N). + Returns: + error_accel (N-2). + """ + # (N-2)x14x3 + accel_gt = gt[:-2] - 2 * gt[1:-1] + gt[2:] + accel_pred = pred[:-2] - 2 * pred[1:-1] + pred[2:] + + normed = np.linalg.norm(accel_pred - accel_gt, axis=2) + + if mask is None: + new_vis = np.ones(len(normed), dtype=bool) + else: + invis = np.logical_not(mask) + invis1 = np.roll(invis, -1) + invis2 = np.roll(invis, -2) + new_invis = np.logical_or(invis, np.logical_or(invis1, invis2))[:-2] + new_vis = np.logical_not(new_invis) + + return np.mean(normed[new_vis], axis=1) + + +def vertice_pve(pred_verts, target_verts): + """Computes per vertex error (PVE). + + Args: + verts_gt (N x verts_num x 3). + verts_pred (N x verts_num x 3). + Returns: + error_verts. + """ + assert len(pred_verts) == len(target_verts) + error = np.linalg.norm(pred_verts - target_verts, ord=2, axis=-1).mean() + return error + + +def keypoint_3d_pck(pred, gt, mask, alignment='none', threshold=150.): + """Calculate the Percentage of Correct Keypoints (3DPCK) w. or w/o rigid + alignment. + Paper ref: `Monocular 3D Human Pose Estimation In The Wild Using Improved + CNN Supervision' 3DV'2017. `__ . + Note: + - batch_size: N + - num_keypoints: K + - keypoint_dims: C + Args: + pred (np.ndarray[N, K, C]): Predicted keypoint location. + gt (np.ndarray[N, K, C]): Groundtruth keypoint location. + mask (np.ndarray[N, K]): Visibility of the target. False for invisible + joints, and True for visible. Invisible joints will be ignored for + accuracy calculation. + alignment (str, optional): method to align the prediction with the + groundtruth. Supported options are: + - ``'none'``: no alignment will be applied + - ``'scale'``: align in the least-square sense in scale + - ``'procrustes'``: align in the least-square sense in scale, + rotation and translation. + threshold: If L2 distance between the prediction and the groundtruth + is less then threshold, the predicted result is considered as + correct. Default: 150 (mm). + Returns: + pck: percentage of correct keypoints. + """ + assert mask.any() + + if alignment == 'none': + pass + elif alignment == 'procrustes': + pred = np.stack([ + compute_similarity_transform(pred_i, gt_i) + for pred_i, gt_i in zip(pred, gt) + ]) + elif alignment == 'scale': + pred_dot_pred = np.einsum('nkc,nkc->n', pred, pred) + pred_dot_gt = np.einsum('nkc,nkc->n', pred, gt) + scale_factor = pred_dot_gt / pred_dot_pred + pred = pred * scale_factor[:, None, None] + else: + raise ValueError(f'Invalid value for alignment: {alignment}') + + error = np.linalg.norm(pred - gt, ord=2, axis=-1) + pck = (error < threshold).astype(np.float32)[mask].mean() * 100 + + return pck + + +def keypoint_3d_auc(pred, gt, mask, alignment='none'): + """Calculate the Area Under the Curve (3DAUC) computed for a range of 3DPCK + thresholds. + Paper ref: `Monocular 3D Human Pose Estimation In The Wild Using Improved + CNN Supervision' 3DV'2017. `__ . + This implementation is derived from mpii_compute_3d_pck.m, which is + provided as part of the MPI-INF-3DHP test data release. + Note: + batch_size: N + num_keypoints: K + keypoint_dims: C + Args: + pred (np.ndarray[N, K, C]): Predicted keypoint location. + gt (np.ndarray[N, K, C]): Groundtruth keypoint location. + mask (np.ndarray[N, K]): Visibility of the target. False for invisible + joints, and True for visible. Invisible joints will be ignored for + accuracy calculation. + alignment (str, optional): method to align the prediction with the + groundtruth. Supported options are: + - ``'none'``: no alignment will be applied + - ``'scale'``: align in the least-square sense in scale + - ``'procrustes'``: align in the least-square sense in scale, + rotation and translation. + Returns: + auc: AUC computed for a range of 3DPCK thresholds. + """ + assert mask.any() + + if alignment == 'none': + pass + elif alignment == 'procrustes': + pred = np.stack([ + compute_similarity_transform(pred_i, gt_i) + for pred_i, gt_i in zip(pred, gt) + ]) + elif alignment == 'scale': + pred_dot_pred = np.einsum('nkc,nkc->n', pred, pred) + pred_dot_gt = np.einsum('nkc,nkc->n', pred, gt) + scale_factor = pred_dot_gt / pred_dot_pred + pred = pred * scale_factor[:, None, None] + else: + raise ValueError(f'Invalid value for alignment: {alignment}') + + error = np.linalg.norm(pred - gt, ord=2, axis=-1) + + thresholds = np.linspace(0., 150, 31) + pck_values = np.zeros(len(thresholds)) + for i in range(len(thresholds)): + pck_values[i] = (error < thresholds[i]).astype(np.float32)[mask].mean() + + auc = pck_values.mean() * 100 + + return auc diff --git a/mmhuman3d/core/evaluation/mpjpe.py b/mmhuman3d/core/evaluation/mpjpe.py deleted file mode 100644 index 5d407420..00000000 --- a/mmhuman3d/core/evaluation/mpjpe.py +++ /dev/null @@ -1,50 +0,0 @@ -# Copyright (c) OpenMMLab. All rights reserved. -import numpy as np - -from .mesh_eval import compute_similarity_transform - - -def keypoint_mpjpe(pred, gt, mask, alignment='none'): - """Calculate the mean per-joint position error (MPJPE) and the error after - rigid alignment with the ground truth (P-MPJPE). - batch_size: N - num_keypoints: K - keypoint_dims: C - Args: - pred (np.ndarray[N, K, C]): Predicted keypoint location. - gt (np.ndarray[N, K, C]): Groundtruth keypoint location. - mask (np.ndarray[N, K]): Visibility of the target. False for invisible - joints, and True for visible. Invisible joints will be ignored for - accuracy calculation. - alignment (str, optional): method to align the prediction with the - groundtruth. Supported options are: - - ``'none'``: no alignment will be applied - - ``'scale'``: align in the least-square sense in scale - - ``'procrustes'``: align in the least-square sense in scale, - rotation and translation. - Returns: - tuple: A tuple containing joint position errors - - mpjpe (float|np.ndarray[N]): mean per-joint position error. - - p-mpjpe (float|np.ndarray[N]): mpjpe after rigid alignment with the - ground truth - """ - assert mask.any() - - if alignment == 'none': - pass - elif alignment == 'procrustes': - pred = np.stack([ - compute_similarity_transform(pred_i, gt_i) - for pred_i, gt_i in zip(pred, gt) - ]) - elif alignment == 'scale': - pred_dot_pred = np.einsum('nkc,nkc->n', pred, pred) - pred_dot_gt = np.einsum('nkc,nkc->n', pred, gt) - scale_factor = pred_dot_gt / pred_dot_pred - pred = pred * scale_factor[:, None, None] - else: - raise ValueError(f'Invalid value for alignment: {alignment}') - - error = np.linalg.norm(pred - gt, ord=2, axis=-1)[mask].mean() - - return error diff --git a/mmhuman3d/data/datasets/human_hybrik_dataset.py b/mmhuman3d/data/datasets/human_hybrik_dataset.py index 56582762..f9f55857 100644 --- a/mmhuman3d/data/datasets/human_hybrik_dataset.py +++ b/mmhuman3d/data/datasets/human_hybrik_dataset.py @@ -3,12 +3,21 @@ import os.path from abc import ABCMeta from collections import OrderedDict +from typing import List, Optional, Union +import mmcv import numpy as np +import torch from mmhuman3d.core.conventions.keypoints_mapping import get_mapping -from mmhuman3d.core.evaluation.mpjpe import keypoint_mpjpe +from mmhuman3d.core.evaluation import ( + keypoint_3d_auc, + keypoint_3d_pck, + keypoint_mpjpe, + vertice_pve, +) from mmhuman3d.data.data_structures.human_data import HumanData +from mmhuman3d.models.builder import build_body_model from mmhuman3d.utils.demo_utils import box2cs, xyxy2xywh from .base_dataset import BaseDataset from .builder import DATASETS @@ -31,11 +40,16 @@ class HybrIKHumanImageDataset(BaseDataset, metaclass=ABCMeta): test_mode (bool): Store True when building test dataset. Default: False. """ + # metric + ALLOWED_METRICS = { + 'mpjpe', 'pa-mpjpe', 'pve', '3dpck', 'pa-3dpck', '3dauc', 'pa-3dauc' + } def __init__(self, data_prefix, pipeline, dataset_name, + body_model, ann_file, test_mode=False): if dataset_name is not None: @@ -43,6 +57,10 @@ def __init__(self, self.test_mode = test_mode super(HybrIKHumanImageDataset, self).__init__(data_prefix, pipeline, ann_file, test_mode) + if body_model is not None: + self.body_model = build_body_model(body_model) + else: + self.body_model = None def get_annotation_file(self): """Obtain annotation file path from data prefix.""" @@ -190,25 +208,71 @@ def load_annotations(self): self.data_infos.append(info) - def evaluate(self, outputs, res_folder, metric='joint_error', logger=None): - """Evaluate 3D keypoint results.""" + def evaluate(self, + outputs: list, + res_folder: str, + metric: Optional[Union[str, List[str]]] = 'pa-mpjpe', + **kwargs: dict): + """Evaluate 3D keypoint results. + + Args: + outputs (list): results from model inference. + res_folder (str): path to store results. + metric (Optional[Union[str, List(str)]]): + the type of metric. Default: 'pa-mpjpe' + kwargs (dict): other arguments. + Returns: + dict: + A dict of all evaluation results. + """ metrics = metric if isinstance(metric, list) else [metric] - allowed_metrics = ['joint_error'] for metric in metrics: - if metric not in allowed_metrics: - raise KeyError(f'metric {metric} is not supported') + if metric not in self.ALLOWED_METRICS: + raise ValueError(f'metric {metric} is not supported') res_file = os.path.join(res_folder, 'result_keypoints.json') - kpts_dict = {} + + res_dict = {} for out in outputs: - for (keypoints, idx) in zip(out['xyz_17'], out['image_idx']): - kpts_dict[int(idx)] = keypoints.tolist() - kpts = [] + target_id = out['image_idx'] + batch_size = len(out['xyz_17']) + for i in range(batch_size): + res_dict[int(target_id[i])] = dict( + keypoints=out['xyz_17'][i], + poses=out['smpl_pose'][i], + betas=out['smpl_beta'][i], + ) + + keypoints, poses, betas = [], [], [] for i in range(self.num_data): - kpts.append(kpts_dict[i]) - self._write_keypoint_results(kpts, res_file) - info_str = self._report_metric(res_file) - name_value = OrderedDict(info_str) + keypoints.append(res_dict[i]['keypoints']) + poses.append(res_dict[i]['poses']) + betas.append(res_dict[i]['betas']) + + res = dict(keypoints=keypoints, poses=poses, betas=betas) + mmcv.dump(res, res_file) + + name_value_tuples = [] + for _metric in metrics: + if _metric == 'mpjpe': + _nv_tuples = self._report_mpjpe(res) + elif _metric == 'pa-mpjpe': + _nv_tuples = self._report_mpjpe(res, metric='pa-mpjpe') + elif _metric == '3dpck': + _nv_tuples = self._report_3d_pck(res) + elif _metric == 'pa-3dpck': + _nv_tuples = self._report_3d_pck(res, metric='pa-3dpck') + elif _metric == '3dauc': + _nv_tuples = self._report_3d_auc(res) + elif _metric == 'pa-3dauc': + _nv_tuples = self._report_3d_auc(res, metric='pa-3dauc') + elif _metric == 'pve': + _nv_tuples = self._report_pve(res) + else: + raise NotImplementedError + name_value_tuples.extend(_nv_tuples) + + name_value = OrderedDict(name_value_tuples) return name_value @staticmethod @@ -217,67 +281,172 @@ def _write_keypoint_results(keypoints, res_file): with open(res_file, 'w') as f: json.dump(keypoints, f, sort_keys=True, indent=4) - def _report_metric(self, res_file): - """Keypoint evaluation. + def _parse_result(self, res, mode='keypoint'): + """Parse results.""" + gts = self.data_infos + if mode == 'vertice': + pred_pose = torch.FloatTensor(res['poses']) + pred_beta = torch.FloatTensor(res['betas']) + pred_output = self.body_model( + betas=pred_beta, + body_pose=pred_pose[:, 1:], + global_orient=pred_pose[:, 0].unsqueeze(1), + pose2rot=False) + pred_vertices = pred_output['vertices'].detach().cpu().numpy() + + gt_pose = torch.FloatTensor([gt['pose'] + for gt in gts]).view(-1, 72) + gt_beta = torch.FloatTensor([gt['beta'] for gt in gts]) + gt_output = self.body_model( + betas=gt_beta, + body_pose=gt_pose[:, 3:], + global_orient=gt_pose[:, :3]) + gt_vertices = gt_output['vertices'].detach().cpu().numpy() + gt_mask = np.ones(gt_vertices.shape[:-1]) + assert len(pred_vertices) == self.num_data + + return pred_vertices * 1000., gt_vertices * 1000., gt_mask + elif mode == 'keypoint': + pred_keypoints3d = res['keypoints'] + assert len(pred_keypoints3d) == self.num_data + # (B, 17, 3) + pred_keypoints3d = np.array(pred_keypoints3d) + factor, root_idx_17 = 1, 0 + + if self.dataset_name == 'mpi_inf_3dhp': + _, hp3d_idxs, _ = get_mapping('human_data', + 'mpi_inf_3dhp_test') + gt_keypoints3d = np.array( + [gt['joint_relative_17'][hp3d_idxs] for gt in gts]) + joint_mapper = [ + 14, 11, 12, 13, 8, 9, 10, 15, 1, 16, 0, 5, 6, 7, 2, 3, 4 + ] + gt_keypoints3d_mask = np.ones( + (len(gt_keypoints3d), len(joint_mapper))) + else: + _, h36m_idxs, _ = get_mapping('human_data', 'h36m') + gt_keypoints3d = np.array( + [gt['joint_relative_17'][h36m_idxs] for gt in gts]) + joint_mapper = [ + 6, 5, 4, 1, 2, 3, 16, 15, 14, 11, 12, 13, 8, 10 + ] + gt_keypoints3d_mask = np.ones( + (len(gt_keypoints3d), len(joint_mapper))) + if self.dataset_name == 'pw3d': + factor = 1000 + + assert len(pred_keypoints3d) == self.num_data + + pred_keypoints3d = pred_keypoints3d * (2000 / factor) + if self.dataset_name == 'mpi_inf_3dhp': + gt_keypoints3d = gt_keypoints3d[:, joint_mapper, :] + # root joint alignment + pred_keypoints3d = ( + pred_keypoints3d - + pred_keypoints3d[:, None, root_idx_17]) * factor + gt_keypoints3d = (gt_keypoints3d - + gt_keypoints3d[:, None, root_idx_17]) * factor + + if self.dataset_name == 'pw3d' or self.dataset_name == 'h36m': + # select eval 14 joints + pred_keypoints3d = pred_keypoints3d[:, joint_mapper, :] + gt_keypoints3d = gt_keypoints3d[:, joint_mapper, :] + + gt_keypoints3d_mask = gt_keypoints3d_mask > 0 + + return pred_keypoints3d, gt_keypoints3d, gt_keypoints3d_mask + + else: + raise NotImplementedError() + + def _report_mpjpe(self, res_file, metric='mpjpe'): + """Cauculate mean per joint position error (MPJPE) or its variants PA- + MPJPE. Report mean per joint position error (MPJPE) and mean per joint - position error after rigid alignment (MPJPE-PA) + position error after rigid alignment (PA-MPJPE) """ + pred_keypoints3d, gt_keypoints3d, gt_keypoints3d_mask = \ + self._parse_result(res_file, mode='keypoint') + + err_name = metric.upper() + if metric == 'mpjpe': + alignment = 'none' + elif metric == 'pa-mpjpe': + alignment = 'procrustes' + else: + raise ValueError(f'Invalid metric: {metric}') - with open(res_file, 'r') as fin: - pred_keypoints3d = json.load(fin) - assert len(pred_keypoints3d) == len(self.data_infos) + error = keypoint_mpjpe(pred_keypoints3d, gt_keypoints3d, + gt_keypoints3d_mask, alignment) + info_str = [(err_name, error)] - pred_keypoints3d = np.array(pred_keypoints3d) - factor, root_idx_17 = 1, 0 - gts = self.data_infos - if self.dataset_name == 'mpi_inf_3dhp': - _, hp3d_idxs, _ = get_mapping('human_data', 'mpi_inf_3dhp_test') - gt_keypoints3d = np.array( - [gt['joint_relative_17'][hp3d_idxs] for gt in gts]) - joint_mapper = [ - 14, 11, 12, 13, 8, 9, 10, 15, 1, 16, 0, 5, 6, 7, 2, 3, 4 - ] - gt_keypoints3d_mask = np.ones( - (len(gt_keypoints3d), len(joint_mapper))) + return info_str + def _report_3d_pck(self, res_file, metric='3dpck'): + """Cauculate Percentage of Correct Keypoints (3DPCK) w. or w/o + Procrustes alignment. + Args: + keypoint_results (list): Keypoint predictions. See + 'Body3DMpiInf3dhpDataset.evaluate' for details. + metric (str): Specify mpjpe variants. Supported options are: + - ``'3dpck'``: Standard 3DPCK. + - ``'pa-3dpck'``: + 3DPCK after aligning prediction to groundtruth + via a rigid transformation (scale, rotation and + translation). + """ + + pred_keypoints3d, gt_keypoints3d, gt_keypoints3d_mask = \ + self._parse_result(res_file, mode='keypoint') + + err_name = metric.upper() + if metric == '3dpck': + alignment = 'none' + elif metric == 'pa-3dpck': + alignment = 'procrustes' else: - _, h36m_idxs, _ = get_mapping('human_data', 'h36m') - gt_keypoints3d = np.array( - [gt['joint_relative_17'][h36m_idxs] for gt in gts]) - joint_mapper = [6, 5, 4, 1, 2, 3, 16, 15, 14, 11, 12, 13, 8, 10] - gt_keypoints3d_mask = np.ones( - (len(gt_keypoints3d), len(joint_mapper))) - if self.dataset_name == 'pw3d': - factor = 1000 - - print('Evaluation start...') - assert len(gts) == len(pred_keypoints3d) - - pred_keypoints3d = pred_keypoints3d * (2000 / factor) - if self.dataset_name == 'mpi_inf_3dhp': - gt_keypoints3d = gt_keypoints3d[:, joint_mapper, :] - # root joint alignment - pred_keypoints3d = pred_keypoints3d - pred_keypoints3d[:, None, - root_idx_17] - gt_keypoints3d = gt_keypoints3d - gt_keypoints3d[:, None, root_idx_17] - - if self.dataset_name == 'pw3d' or self.dataset_name == 'h36m': - # select eval 14 joints - pred_keypoints3d = pred_keypoints3d[:, joint_mapper, :] - gt_keypoints3d = gt_keypoints3d[:, joint_mapper, :] - - gt_keypoints3d_mask = gt_keypoints3d_mask > 0 - - mpjpe = keypoint_mpjpe(pred_keypoints3d, gt_keypoints3d, - gt_keypoints3d_mask) - mpjpe_pa = keypoint_mpjpe( - pred_keypoints3d, - gt_keypoints3d, - gt_keypoints3d_mask, - alignment='procrustes') - - info_str = [] - info_str.append(('MPJPE', mpjpe * factor)) - info_str.append(('MPJPE-PA', mpjpe_pa * factor)) - return info_str + raise ValueError(f'Invalid metric: {metric}') + + error = keypoint_3d_pck(pred_keypoints3d, gt_keypoints3d, + gt_keypoints3d_mask, alignment) + name_value_tuples = [(err_name, error)] + + return name_value_tuples + + def _report_3d_auc(self, res_file, metric='3dauc'): + """Cauculate the Area Under the Curve (AUC) computed for a range of + 3DPCK thresholds. + Args: + keypoint_results (list): Keypoint predictions. See + 'Body3DMpiInf3dhpDataset.evaluate' for details. + metric (str): Specify mpjpe variants. Supported options are: + - ``'3dauc'``: Standard 3DAUC. + - ``'pa-3dauc'``: 3DAUC after aligning prediction to + groundtruth via a rigid transformation (scale, rotation and + translation). + """ + + pred_keypoints3d, gt_keypoints3d, gt_keypoints3d_mask = \ + self._parse_result(res_file, mode='keypoint') + + err_name = metric.upper() + if metric == '3dauc': + alignment = 'none' + elif metric == 'pa-3dauc': + alignment = 'procrustes' + else: + raise ValueError(f'Invalid metric: {metric}') + + error = keypoint_3d_auc(pred_keypoints3d, gt_keypoints3d, + gt_keypoints3d_mask, alignment) + name_value_tuples = [(err_name, error)] + + return name_value_tuples + + def _report_pve(self, res_file): + """Cauculate per vertex error.""" + pred_verts, gt_verts, _ = \ + self._parse_result(res_file, mode='vertice') + error = vertice_pve(pred_verts, gt_verts) + return [('PVE', error)] diff --git a/mmhuman3d/data/datasets/human_image_dataset.py b/mmhuman3d/data/datasets/human_image_dataset.py index adab52b6..bdf7b35c 100644 --- a/mmhuman3d/data/datasets/human_image_dataset.py +++ b/mmhuman3d/data/datasets/human_image_dataset.py @@ -3,8 +3,9 @@ import os.path from abc import ABCMeta from collections import OrderedDict -from typing import Any, Optional, Union +from typing import Any, List, Optional, Union +import mmcv import numpy as np import torch @@ -12,7 +13,12 @@ convert_kps, get_keypoint_num, ) -from mmhuman3d.core.evaluation.mpjpe import keypoint_mpjpe +from mmhuman3d.core.evaluation import ( + keypoint_3d_auc, + keypoint_3d_pck, + keypoint_mpjpe, + vertice_pve, +) from mmhuman3d.data.data_structures.human_data import HumanData from mmhuman3d.models.builder import build_body_model from .base_dataset import BaseDataset @@ -42,6 +48,10 @@ class HumanImageDataset(BaseDataset, metaclass=ABCMeta): test_mode (bool, optional): in train mode or test mode. Default: False. """ + # metric + ALLOWED_METRICS = { + 'mpjpe', 'pa-mpjpe', 'pve', '3dpck', 'pa-3dpck', '3dauc', 'pa-3dauc' + } def __init__(self, data_prefix: str, @@ -182,36 +192,69 @@ def prepare_data(self, idx: int): def evaluate(self, outputs: list, res_folder: str, - metric: Optional[str] = 'joint_error'): + metric: Optional[Union[str, List[str]]] = 'pa-mpjpe', + **kwargs: dict): """Evaluate 3D keypoint results. Args: outputs (list): results from model inference. res_folder (str): path to store results. - metric (str): the type of metric. Default: 'joint_error' - + metric (Optional[Union[str, List(str)]]): + the type of metric. Default: 'pa-mpjpe' + kwargs (dict): other arguments. Returns: dict: A dict of all evaluation results. """ metrics = metric if isinstance(metric, list) else [metric] - allowed_metrics = ['joint_error'] for metric in metrics: - if metric not in allowed_metrics: + if metric not in self.ALLOWED_METRICS: raise KeyError(f'metric {metric} is not supported') res_file = os.path.join(res_folder, 'result_keypoints.json') # for keeping correctness during multi-gpu test, we sort all results - kpts_dict = {} + + res_dict = {} for out in outputs: - for (keypoints, idx) in zip(out['keypoints_3d'], out['image_idx']): - kpts_dict[int(idx)] = keypoints.tolist() - kpts = [] + target_id = out['image_idx'] + batch_size = len(out['keypoints_3d']) + for i in range(batch_size): + res_dict[int(target_id[i])] = dict( + keypoints=out['keypoints_3d'][i], + poses=out['smpl_pose'][i], + betas=out['smpl_beta'][i], + ) + + keypoints, poses, betas = [], [], [] for i in range(self.num_data): - kpts.append(kpts_dict[i]) - self._write_keypoint_results(kpts, res_file) - info_str = self._report_metric(res_file) - name_value = OrderedDict(info_str) + keypoints.append(res_dict[i]['keypoints']) + poses.append(res_dict[i]['poses']) + betas.append(res_dict[i]['betas']) + + res = dict(keypoints=keypoints, poses=poses, betas=betas) + mmcv.dump(res, res_file) + + name_value_tuples = [] + for _metric in metrics: + if _metric == 'mpjpe': + _nv_tuples = self._report_mpjpe(res) + elif _metric == 'pa-mpjpe': + _nv_tuples = self._report_mpjpe(res, metric='pa-mpjpe') + elif _metric == '3dpck': + _nv_tuples = self._report_3d_pck(res) + elif _metric == 'pa-3dpck': + _nv_tuples = self._report_3d_pck(res, metric='pa-3dpck') + elif _metric == '3dauc': + _nv_tuples = self._report_3d_auc(res) + elif _metric == 'pa-3dauc': + _nv_tuples = self._report_3d_auc(res, metric='pa-3dauc') + elif _metric == 'pve': + _nv_tuples = self._report_pve(res) + else: + raise NotImplementedError + name_value_tuples.extend(_nv_tuples) + + name_value = OrderedDict(name_value_tuples) return name_value @staticmethod @@ -221,150 +264,229 @@ def _write_keypoint_results(keypoints: Any, res_file: str): with open(res_file, 'w') as f: json.dump(keypoints, f, sort_keys=True, indent=4) - def _report_metric(self, res_file: str): - """Keypoint evaluation. + def _parse_result(self, res, mode='keypoint'): + """Parse results.""" - Report mean per joint position error (MPJPE) and mean per joint - position error after rigid alignment (MPJPE-PA) - """ - - with open(res_file, 'r') as fin: - pred_keypoints3d = json.load(fin) - assert len(pred_keypoints3d) == self.num_data - - pred_keypoints3d = np.array(pred_keypoints3d) - if self.dataset_name == 'pw3d': - betas = [] - body_pose = [] - global_orient = [] - gender = [] - smpl_dict = self.human_data['smpl'] + if mode == 'vertice': + # gt + gt_beta, gt_pose, gt_global_orient, gender = [], [], [], [] + gt_smpl_dict = self.human_data['smpl'] for idx in range(self.num_data): - betas.append(smpl_dict['betas'][idx]) - body_pose.append(smpl_dict['body_pose'][idx]) - global_orient.append(smpl_dict['global_orient'][idx]) + gt_beta.append(gt_smpl_dict['betas'][idx]) + gt_pose.append(gt_smpl_dict['body_pose'][idx]) + gt_global_orient.append(gt_smpl_dict['global_orient'][idx]) if self.human_data['meta']['gender'][idx] == 'm': gender.append(0) else: gender.append(1) - betas = torch.FloatTensor(betas) - body_pose = torch.FloatTensor(body_pose).view(-1, 69) - global_orient = torch.FloatTensor(global_orient) + gt_beta = torch.FloatTensor(gt_beta) + gt_pose = torch.FloatTensor(gt_pose).view(-1, 69) + gt_global_orient = torch.FloatTensor(gt_global_orient) gender = torch.Tensor(gender) gt_output = self.body_model( - betas=betas, - body_pose=body_pose, - global_orient=global_orient, + betas=gt_beta, + body_pose=gt_pose, + global_orient=gt_global_orient, gender=gender) - gt_keypoints3d = gt_output['joints'].detach().cpu().numpy() - gt_keypoints3d_mask = np.ones((len(pred_keypoints3d), 24)) - elif self.dataset_name == 'humman': - betas = [] - body_pose = [] - global_orient = [] - smpl_dict = self.human_data['smpl'] - for idx in range(self.num_data): - betas.append(smpl_dict['betas'][idx]) - body_pose.append(smpl_dict['body_pose'][idx]) - global_orient.append(smpl_dict['global_orient'][idx]) - betas = torch.FloatTensor(betas) - body_pose = torch.FloatTensor(body_pose).view(-1, 69) - global_orient = torch.FloatTensor(global_orient) - gt_output = self.body_model( - betas=betas, body_pose=body_pose, global_orient=global_orient) - gt_keypoints3d = gt_output['joints'].detach().cpu().numpy() - gt_keypoints3d_mask = np.ones((len(pred_keypoints3d), 24)) - elif self.dataset_name == 'h36m': - gt_keypoints3d = self.human_data['keypoints3d'][:, :, :3] - gt_keypoints3d_mask = np.ones((len(pred_keypoints3d), 17)) - else: - raise NotImplementedError() + gt_vertices = gt_output['vertices'].detach().cpu().numpy() * 1000. + gt_mask = np.ones(gt_vertices.shape[:-1]) + # pred + pred_pose = torch.FloatTensor(res['poses']) + pred_beta = torch.FloatTensor(res['betas']) + pred_output = self.body_model( + betas=pred_beta, + body_pose=pred_pose[:, 1:], + global_orient=pred_pose[:, 0].unsqueeze(1), + pose2rot=False, + gender=gender) + pred_vertices = pred_output['vertices'].detach().cpu().numpy( + ) * 1000. + + assert len(pred_vertices) == self.num_data + + return pred_vertices, gt_vertices, gt_mask + elif mode == 'keypoint': + pred_keypoints3d = res['keypoints'] + assert len(pred_keypoints3d) == self.num_data + # (B, 17, 3) + pred_keypoints3d = np.array(pred_keypoints3d) + + if self.dataset_name == 'pw3d': + betas = [] + body_pose = [] + global_orient = [] + gender = [] + smpl_dict = self.human_data['smpl'] + for idx in range(self.num_data): + betas.append(smpl_dict['betas'][idx]) + body_pose.append(smpl_dict['body_pose'][idx]) + global_orient.append(smpl_dict['global_orient'][idx]) + if self.human_data['meta']['gender'][idx] == 'm': + gender.append(0) + else: + gender.append(1) + betas = torch.FloatTensor(betas) + body_pose = torch.FloatTensor(body_pose).view(-1, 69) + global_orient = torch.FloatTensor(global_orient) + gender = torch.Tensor(gender) + gt_output = self.body_model( + betas=betas, + body_pose=body_pose, + global_orient=global_orient, + gender=gender) + gt_keypoints3d = gt_output['joints'].detach().cpu().numpy() + gt_keypoints3d_mask = np.ones((len(pred_keypoints3d), 24)) + elif self.dataset_name in ['h36m', 'humman']: + gt_keypoints3d = self.human_data['keypoints3d'][:, :, :3] + gt_keypoints3d_mask = np.ones((len(pred_keypoints3d), 17)) - # SMPL_49 only! - if gt_keypoints3d.shape[1] == 49: - assert pred_keypoints3d.shape[1] == 49 + else: + raise NotImplementedError() - gt_keypoints3d = gt_keypoints3d[:, 25:, :] - pred_keypoints3d = pred_keypoints3d[:, 25:, :] + # SMPL_49 only! + if gt_keypoints3d.shape[1] == 49: + assert pred_keypoints3d.shape[1] == 49 - joint_mapper = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 18] - gt_keypoints3d = gt_keypoints3d[:, joint_mapper, :] - pred_keypoints3d = pred_keypoints3d[:, joint_mapper, :] + gt_keypoints3d = gt_keypoints3d[:, 25:, :] + pred_keypoints3d = pred_keypoints3d[:, 25:, :] - # we only evaluate on 14 lsp joints - pred_pelvis = (pred_keypoints3d[:, 2] + pred_keypoints3d[:, 3]) / 2 - gt_pelvis = (gt_keypoints3d[:, 2] + gt_keypoints3d[:, 3]) / 2 + joint_mapper = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 18] + gt_keypoints3d = gt_keypoints3d[:, joint_mapper, :] + pred_keypoints3d = pred_keypoints3d[:, joint_mapper, :] - # H36M for testing! - elif gt_keypoints3d.shape[1] == 17: - assert pred_keypoints3d.shape[1] == 17 + # we only evaluate on 14 lsp joints + pred_pelvis = (pred_keypoints3d[:, 2] + + pred_keypoints3d[:, 3]) / 2 + gt_pelvis = (gt_keypoints3d[:, 2] + gt_keypoints3d[:, 3]) / 2 - H36M_TO_J17 = [ - 6, 5, 4, 1, 2, 3, 16, 15, 14, 11, 12, 13, 8, 10, 0, 7, 9 - ] - H36M_TO_J14 = H36M_TO_J17[:14] - joint_mapper = H36M_TO_J14 + # H36M for testing! + elif gt_keypoints3d.shape[1] == 17: + assert pred_keypoints3d.shape[1] == 17 - pred_pelvis = pred_keypoints3d[:, 0] - gt_pelvis = gt_keypoints3d[:, 0] + H36M_TO_J17 = [ + 6, 5, 4, 1, 2, 3, 16, 15, 14, 11, 12, 13, 8, 10, 0, 7, 9 + ] + H36M_TO_J14 = H36M_TO_J17[:14] + joint_mapper = H36M_TO_J14 - gt_keypoints3d = gt_keypoints3d[:, joint_mapper, :] - pred_keypoints3d = pred_keypoints3d[:, joint_mapper, :] + pred_pelvis = pred_keypoints3d[:, 0] + gt_pelvis = gt_keypoints3d[:, 0] - # keypoint 24 - elif gt_keypoints3d.shape[1] == 24: - assert pred_keypoints3d.shape[1] == 24 + gt_keypoints3d = gt_keypoints3d[:, joint_mapper, :] + pred_keypoints3d = pred_keypoints3d[:, joint_mapper, :] - joint_mapper = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 18] - gt_keypoints3d = gt_keypoints3d[:, joint_mapper, :] - pred_keypoints3d = pred_keypoints3d[:, joint_mapper, :] + # keypoint 24 + elif gt_keypoints3d.shape[1] == 24: + assert pred_keypoints3d.shape[1] == 24 - # we only evaluate on 14 lsp joints - pred_pelvis = (pred_keypoints3d[:, 2] + pred_keypoints3d[:, 3]) / 2 - gt_pelvis = (gt_keypoints3d[:, 2] + gt_keypoints3d[:, 3]) / 2 + joint_mapper = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 18] + gt_keypoints3d = gt_keypoints3d[:, joint_mapper, :] + pred_keypoints3d = pred_keypoints3d[:, joint_mapper, :] - # humman keypoints (not SMPL keypoints) - elif gt_keypoints3d.shape[1] == 133: - assert pred_keypoints3d.shape[1] == 17 + # we only evaluate on 14 lsp joints + pred_pelvis = (pred_keypoints3d[:, 2] + + pred_keypoints3d[:, 3]) / 2 + gt_pelvis = (gt_keypoints3d[:, 2] + gt_keypoints3d[:, 3]) / 2 - H36M_TO_J17 = [ - 6, 5, 4, 1, 2, 3, 16, 15, 14, 11, 12, 13, 8, 10, 0, 7, 9 - ] - H36M_TO_J14 = H36M_TO_J17[:14] - pred_joint_mapper = H36M_TO_J14 - pred_keypoints3d = pred_keypoints3d[:, pred_joint_mapper, :] + else: + pass - # the last two are not mapped - gt_joint_mapper = [16, 14, 12, 11, 13, 15, 10, 8, 6, 5, 7, 9, 0, 0] - gt_keypoints3d = gt_keypoints3d[:, gt_joint_mapper, :] + pred_keypoints3d = (pred_keypoints3d - + pred_pelvis[:, None, :]) * 1000 + gt_keypoints3d = (gt_keypoints3d - gt_pelvis[:, None, :]) * 1000 - pred_pelvis = (pred_keypoints3d[:, 2] + pred_keypoints3d[:, 3]) / 2 - gt_pelvis = (gt_keypoints3d[:, 2] + gt_keypoints3d[:, 3]) / 2 + gt_keypoints3d_mask = gt_keypoints3d_mask[:, joint_mapper] > 0 - # TODO: temp solution - joint_mapper = None - gt_keypoints3d_mask = np.ones((len(pred_keypoints3d), 14)) - gt_keypoints3d_mask[:, 12:14] = 0 # the last two are invalid - gt_keypoints3d_mask = gt_keypoints3d_mask > 0 + return pred_keypoints3d, gt_keypoints3d, gt_keypoints3d_mask - else: - raise NotImplementedError + def _report_mpjpe(self, res_file, metric='mpjpe'): + """Cauculate mean per joint position error (MPJPE) or its variants PA- + MPJPE. - pred_keypoints3d = pred_keypoints3d - pred_pelvis[:, None, :] - gt_keypoints3d = gt_keypoints3d - gt_pelvis[:, None, :] + Report mean per joint position error (MPJPE) and mean per joint + position error after rigid alignment (PA-MPJPE) + """ + pred_keypoints3d, gt_keypoints3d, gt_keypoints3d_mask = \ + self._parse_result(res_file, mode='keypoint') + + err_name = metric.upper() + if metric == 'mpjpe': + alignment = 'none' + elif metric == 'pa-mpjpe': + alignment = 'procrustes' + else: + raise ValueError(f'Invalid metric: {metric}') - if joint_mapper is not None: - gt_keypoints3d_mask = gt_keypoints3d_mask[:, joint_mapper] > 0 + error = keypoint_mpjpe(pred_keypoints3d, gt_keypoints3d, + gt_keypoints3d_mask, alignment) + info_str = [(err_name, error)] - mpjpe = keypoint_mpjpe(pred_keypoints3d, gt_keypoints3d, - gt_keypoints3d_mask) - mpjpe_pa = keypoint_mpjpe( - pred_keypoints3d, - gt_keypoints3d, - gt_keypoints3d_mask, - alignment='procrustes') - - info_str = [] - info_str.append(('MPJPE', mpjpe * 1000)) - info_str.append(('MPJPE-PA', mpjpe_pa * 1000)) return info_str + + def _report_3d_pck(self, res_file, metric='3dpck'): + """Cauculate Percentage of Correct Keypoints (3DPCK) w. or w/o + Procrustes alignment. + Args: + keypoint_results (list): Keypoint predictions. See + 'Body3DMpiInf3dhpDataset.evaluate' for details. + metric (str): Specify mpjpe variants. Supported options are: + - ``'3dpck'``: Standard 3DPCK. + - ``'pa-3dpck'``: + 3DPCK after aligning prediction to groundtruth + via a rigid transformation (scale, rotation and + translation). + """ + + pred_keypoints3d, gt_keypoints3d, gt_keypoints3d_mask = \ + self._parse_result(res_file) + + err_name = metric.upper() + if metric == '3dpck': + alignment = 'none' + elif metric == 'pa-3dpck': + alignment = 'procrustes' + else: + raise ValueError(f'Invalid metric: {metric}') + + error = keypoint_3d_pck(pred_keypoints3d, gt_keypoints3d, + gt_keypoints3d_mask, alignment) + name_value_tuples = [(err_name, error)] + + return name_value_tuples + + def _report_3d_auc(self, res_file, metric='3dauc'): + """Cauculate the Area Under the Curve (AUC) computed for a range of + 3DPCK thresholds. + Args: + keypoint_results (list): Keypoint predictions. See + 'Body3DMpiInf3dhpDataset.evaluate' for details. + metric (str): Specify mpjpe variants. Supported options are: + - ``'3dauc'``: Standard 3DAUC. + - ``'pa-3dauc'``: 3DAUC after aligning prediction to + groundtruth via a rigid transformation (scale, rotation and + translation). + """ + + pred_keypoints3d, gt_keypoints3d, gt_keypoints3d_mask = \ + self._parse_result(res_file) + + err_name = metric.upper() + if metric == '3dauc': + alignment = 'none' + elif metric == 'pa-3dauc': + alignment = 'procrustes' + else: + raise ValueError(f'Invalid metric: {metric}') + + error = keypoint_3d_auc(pred_keypoints3d, gt_keypoints3d, + gt_keypoints3d_mask, alignment) + name_value_tuples = [(err_name, error)] + + return name_value_tuples + + def _report_pve(self, res_file): + """Cauculate per vertex error.""" + pred_verts, gt_verts, _ = \ + self._parse_result(res_file, mode='vertice') + error = vertice_pve(pred_verts, gt_verts) + return [('PVE', error)] diff --git a/mmhuman3d/models/__init__.py b/mmhuman3d/models/__init__.py index 2be88b52..4ae52f2b 100644 --- a/mmhuman3d/models/__init__.py +++ b/mmhuman3d/models/__init__.py @@ -16,6 +16,7 @@ build_head, build_loss, build_neck, + build_registrant, ) from .discriminators import * # noqa: F401,F403 from .heads import * # noqa: F401,F403 @@ -26,5 +27,6 @@ __all__ = [ 'BACKBONES', 'LOSSES', 'ARCHITECTURES', 'HEADS', 'BODY_MODELS', 'NECKS', 'DISCRIMINATORS', 'build_backbone', 'build_loss', 'build_architecture', - 'build_body_model', 'build_head', 'build_neck', 'build_discriminator' + 'build_body_model', 'build_head', 'build_neck', 'build_discriminator', + 'build_registrant' ] diff --git a/mmhuman3d/models/architectures/hybrik.py b/mmhuman3d/models/architectures/hybrik.py index 58e22d13..40e33ddc 100644 --- a/mmhuman3d/models/architectures/hybrik.py +++ b/mmhuman3d/models/architectures/hybrik.py @@ -241,6 +241,8 @@ def forward_test(self, img, img_metas, **kwargs): pred_xyz_jts_17 = pred_xyz_jts_17.cpu().data.numpy() pred_uvd_jts = pred_uvd_jts.cpu().data pred_mesh = pred_mesh.cpu().data.numpy() + pred_pose = output['pred_pose'].cpu().data.numpy() + pred_beta = output['pred_shape'].cpu().data.numpy() assert pred_xyz_jts_17.ndim in [2, 3] pred_xyz_jts_17 = pred_xyz_jts_17.reshape(pred_xyz_jts_17.shape[0], 17, @@ -264,6 +266,8 @@ def forward_test(self, img, img_metas, **kwargs): all_preds = {} all_preds['vertices'] = pred_mesh + all_preds['smpl_pose'] = pred_pose + all_preds['smpl_beta'] = pred_beta all_preds['xyz_17'] = pred_xyz_jts_17 all_preds['uvd_jts'] = pose_coords all_preds['xyz_24'] = pred_xyz_jts_24_struct diff --git a/mmhuman3d/models/body_models/smpl.py b/mmhuman3d/models/body_models/smpl.py index 98a8d9d5..f52f89fe 100644 --- a/mmhuman3d/models/body_models/smpl.py +++ b/mmhuman3d/models/body_models/smpl.py @@ -13,7 +13,7 @@ ) from mmhuman3d.core.conventions.segmentation import body_segmentation from mmhuman3d.models.utils import batch_inverse_kinematics_transform -from mmhuman3d.utils.transforms import quat_to_rotmat, rotmat_to_quat +from mmhuman3d.utils.transforms import quat_to_rotmat from ..builder import BODY_MODELS @@ -596,27 +596,21 @@ def forward(self, joints_from_verts = vertices2joints(self.joints_regressor_extra, vertices) - rot_mats = rot_mats.reshape(batch_size * 24, 3, 3) - rot_mats = rotmat_to_quat(rot_mats).reshape(batch_size, 24 * 4) - + # rot_mats = rot_mats.reshape(batch_size * 24, 3, 3) if transl is not None: new_joints += transl.unsqueeze(dim=1) vertices += transl.unsqueeze(dim=1) joints_from_verts += transl.unsqueeze(dim=1) else: - vertices = vertices - joints_from_verts[:, self. - root_idx_17, :].unsqueeze( - 1).detach() - new_joints = new_joints - new_joints[:, self. - root_idx_smpl, :].unsqueeze( - 1).detach() + new_joints = new_joints - \ + new_joints[:, self.root_idx_smpl, :].unsqueeze(1).detach() joints_from_verts = joints_from_verts - \ joints_from_verts[:, self.root_idx_17, :].unsqueeze(1).detach() output = { 'vertices': vertices, 'joints': new_joints, - 'rot_mats': rot_mats, + 'poses': rot_mats, 'joints_from_verts': joints_from_verts, } return output diff --git a/mmhuman3d/models/heads/hybrik_head.py b/mmhuman3d/models/heads/hybrik_head.py index 4aedc7eb..a8f10e14 100644 --- a/mmhuman3d/models/heads/hybrik_head.py +++ b/mmhuman3d/models/heads/hybrik_head.py @@ -427,8 +427,8 @@ def forward(self, pred_xyz_jts_24_struct = hybrik_output['joints'].float() / 2 # -0.5 ~ 0.5 pred_xyz_jts_17 = hybrik_output['joints_from_verts'].float() / 2 - pred_theta_mats = hybrik_output['rot_mats'].float().reshape( - batch_size, 24 * 4) + pred_poses = hybrik_output['poses'].float().reshape( + batch_size, 24, 3, 3) pred_xyz_jts_24 = pred_xyz_jts_29[:, :24, :].reshape(batch_size, 72) pred_xyz_jts_24_struct = pred_xyz_jts_24_struct.reshape(batch_size, 72) pred_xyz_jts_17 = pred_xyz_jts_17.reshape(batch_size, 17 * 3) @@ -437,7 +437,7 @@ def forward(self, 'pred_phi': pred_phi, 'pred_delta_shape': delta_shape, 'pred_shape': pred_shape, - 'pred_theta_mats': pred_theta_mats, + 'pred_pose': pred_poses, 'pred_uvd_jts': pred_uvd_jts_29_flat, 'pred_xyz_jts_24': pred_xyz_jts_24, 'pred_xyz_jts_24_struct': pred_xyz_jts_24_struct, diff --git a/tests/test_datasets/test_human_image_dataset.py b/tests/test_datasets/test_human_image_dataset.py index a4e0c5d6..f5c63a26 100644 --- a/tests/test_datasets/test_human_image_dataset.py +++ b/tests/test_datasets/test_human_image_dataset.py @@ -36,16 +36,45 @@ def test_human_image_dataset(): keypoint_dst='h36m', model_path='data/body_models/smpl'), ann_file='sample_3dpw_test.npz') - test_dataset.num_data = 1 + test_dataset.num_data = num_data outputs = [{ 'keypoints_3d': np.random.rand(num_data, 17, 3), + 'smpl_pose': np.random.rand(num_data, 24, 3, 3), + 'smpl_beta': np.random.rand(num_data, 10), 'image_idx': np.arange(num_data) }] res = test_dataset.evaluate(outputs, res_folder='tests/data') + assert 'PA-MPJPE' in res + assert res['PA-MPJPE'] > 0 + + res = test_dataset.evaluate( + outputs, res_folder='tests/data', metric='mpjpe') assert 'MPJPE' in res - assert 'MPJPE-PA' in res assert res['MPJPE'] > 0 - assert res['MPJPE-PA'] > 0 + + res = test_dataset.evaluate(outputs, res_folder='tests/data', metric='pve') + assert 'PVE' in res + assert res['PVE'] > 0 + + res = test_dataset.evaluate( + outputs, res_folder='tests/data', metric='pa-3dpck') + assert 'PA-3DPCK' in res + assert res['PA-3DPCK'] >= 0 + + res = test_dataset.evaluate( + outputs, res_folder='tests/data', metric='3dpck') + assert '3DPCK' in res + assert res['3DPCK'] >= 0 + + res = test_dataset.evaluate( + outputs, res_folder='tests/data', metric='pa-3dauc') + assert 'PA-3DAUC' in res + assert res['PA-3DAUC'] >= 0 + + res = test_dataset.evaluate( + outputs, res_folder='tests/data', metric='3dauc') + assert '3DAUC' in res + assert res['3DAUC'] >= 0 test_dataset = HumanImageDataset( data_prefix='tests/data', @@ -60,13 +89,13 @@ def test_human_image_dataset(): test_dataset.num_data = 1 outputs = [{ 'keypoints_3d': np.random.rand(num_data, 24, 3), + 'smpl_pose': np.random.rand(num_data, 24, 3, 3), + 'smpl_beta': np.random.rand(num_data, 10), 'image_idx': np.arange(num_data) }] res = test_dataset.evaluate(outputs, res_folder='tests/data') - assert 'MPJPE' in res - assert 'MPJPE-PA' in res - assert res['MPJPE'] > 0 - assert res['MPJPE-PA'] > 0 + assert 'PA-MPJPE' in res + assert res['PA-MPJPE'] > 0 test_dataset = HumanImageDataset( data_prefix='tests/data', @@ -81,13 +110,13 @@ def test_human_image_dataset(): test_dataset.num_data = 1 outputs = [{ 'keypoints_3d': np.random.rand(num_data, 49, 3), + 'smpl_pose': np.random.rand(num_data, 24, 3, 3), + 'smpl_beta': np.random.rand(num_data, 10), 'image_idx': np.arange(num_data) }] res = test_dataset.evaluate(outputs, res_folder='tests/data') - assert 'MPJPE' in res - assert 'MPJPE-PA' in res - assert res['MPJPE'] > 0 - assert res['MPJPE-PA'] > 0 + assert 'PA-MPJPE' in res + assert res['PA-MPJPE'] > 0 def test_pipeline(): diff --git a/tests/test_eval_hook.py b/tests/test_eval_hook.py new file mode 100644 index 00000000..6dabbc4b --- /dev/null +++ b/tests/test_eval_hook.py @@ -0,0 +1,289 @@ +# Copyright (c) OpenMMLab. All rights reserved. +import os.path as osp +import tempfile +import unittest.mock as mock +from collections import OrderedDict +from unittest.mock import MagicMock, patch + +import pytest +import torch +import torch.nn as nn +from mmcv.runner import EpochBasedRunner, build_optimizer +from mmcv.utils import get_logger +from torch.utils.data import DataLoader, Dataset + +from mmhuman3d.core.evaluation import DistEvalHook, EvalHook + + +class ExampleDataset(Dataset): + + def __init__(self): + self.index = 0 + self.eval_result = [0.1, 0.4, 0.3, 0.05, 0.2, 0.7, 0.4, 0.6] + + def __getitem__(self, idx): + results = dict(imgs=torch.tensor([1])) + return results + + def __len__(self): + return 1 + + @mock.create_autospec + def evaluate(self, results, res_folder=None, logger=None): + pass + + +class EvalDataset(ExampleDataset): + + def evaluate(self, results, res_folder=None, logger=None): + acc = self.eval_result[self.index] + output = OrderedDict(mpjpe=acc, index=self.index, pve=acc) + self.index += 1 + return output + + +class ExampleModel(nn.Module): + + def __init__(self): + super().__init__() + self.conv = nn.Linear(1, 1) + self.test_cfg = None + + def forward(self, imgs, return_loss=False): + return imgs + + def train_step(self, data_batch, optimizer, **kwargs): + outputs = { + 'loss': 0.5, + 'log_vars': { + 'accuracy': 0.98 + }, + 'num_samples': 1 + } + return outputs + + +@patch('mmhuman3d.apis.single_gpu_test', MagicMock) +@patch('mmhuman3d.apis.multi_gpu_test', MagicMock) +@pytest.mark.parametrize('EvalHookCls', (EvalHook, DistEvalHook)) +def test_eval_hook(EvalHookCls): + with pytest.raises(TypeError): + # dataloader must be a pytorch DataLoader + test_dataset = ExampleDataset() + data_loader = [ + DataLoader( + test_dataset, + batch_size=1, + sampler=None, + num_workers=0, + shuffle=False) + ] + EvalHookCls(data_loader) + + with pytest.raises(ValueError): + test_dataset = ExampleDataset() + data_loader = DataLoader( + test_dataset, + batch_size=1, + sampler=None, + num_workers=0, + shuffle=False) + # key_indicator should not be None, + # when save_best is set to True + EvalHookCls(data_loader, save_best=True, gpu_collect=False) + + with pytest.raises(KeyError): + # rule must be in keys of rule_map + test_dataset = ExampleDataset() + data_loader = DataLoader( + test_dataset, + batch_size=1, + sampler=None, + num_workers=0, + shuffle=False) + EvalHookCls(data_loader, save_best='auto', rule='unsupport') + + with pytest.raises(ValueError): + # save_best must be valid when rule_map is None + test_dataset = ExampleDataset() + data_loader = DataLoader( + test_dataset, + batch_size=1, + sampler=None, + num_workers=0, + shuffle=False) + EvalHookCls(data_loader, save_best='unsupport') + + optimizer_cfg = dict( + type='SGD', lr=0.01, momentum=0.9, weight_decay=0.0001) + + test_dataset = ExampleDataset() + loader = DataLoader(test_dataset, batch_size=1) + model = ExampleModel() + optimizer = build_optimizer(model, optimizer_cfg) + + data_loader = DataLoader(test_dataset, batch_size=1) + eval_hook = EvalHookCls(data_loader, save_best=None) + with tempfile.TemporaryDirectory() as tmpdir: + logger = get_logger('test_eval') + runner = EpochBasedRunner( + model=model, + batch_processor=None, + optimizer=optimizer, + work_dir=tmpdir, + logger=logger, + max_epochs=1) + runner.register_hook(eval_hook) + runner.run([loader], [('train', 1)]) + assert runner.meta is None or 'best_score' not in runner.meta[ + 'hook_msgs'] + assert runner.meta is None or 'best_ckpt' not in runner.meta[ + 'hook_msgs'] + + # when `save_best` is set to 'auto', first metric will be used. + loader = DataLoader(EvalDataset(), batch_size=1) + model = ExampleModel() + data_loader = DataLoader(EvalDataset(), batch_size=1) + eval_hook = EvalHookCls(data_loader, interval=1, save_best='auto') + + with tempfile.TemporaryDirectory() as tmpdir: + logger = get_logger('test_eval') + runner = EpochBasedRunner( + model=model, + batch_processor=None, + optimizer=optimizer, + work_dir=tmpdir, + logger=logger, + max_epochs=8) + runner.register_checkpoint_hook(dict(interval=1)) + runner.register_hook(eval_hook) + runner.run([loader], [('train', 1)]) + + real_path = osp.join(tmpdir, 'best_mpjpe_epoch_4.pth') + + assert runner.meta['hook_msgs']['best_ckpt'] == osp.realpath(real_path) + assert runner.meta['hook_msgs']['best_score'] == 0.05 + + loader = DataLoader(EvalDataset(), batch_size=1) + model = ExampleModel() + data_loader = DataLoader(EvalDataset(), batch_size=1) + eval_hook = EvalHookCls(data_loader, interval=1, save_best='mpjpe') + + with tempfile.TemporaryDirectory() as tmpdir: + logger = get_logger('test_eval') + runner = EpochBasedRunner( + model=model, + batch_processor=None, + optimizer=optimizer, + work_dir=tmpdir, + logger=logger, + max_epochs=8) + runner.register_checkpoint_hook(dict(interval=1)) + runner.register_hook(eval_hook) + runner.run([loader], [('train', 1)]) + + real_path = osp.join(tmpdir, 'best_mpjpe_epoch_4.pth') + + assert runner.meta['hook_msgs']['best_ckpt'] == osp.realpath(real_path) + assert runner.meta['hook_msgs']['best_score'] == 0.05 + + # update "save_best" according to "key_indicator" + data_loader = DataLoader(EvalDataset(), batch_size=1) + eval_hook = EvalHookCls( + data_loader, key_indicator='mpjpe', save_best=True, rule='less') + with tempfile.TemporaryDirectory() as tmpdir: + logger = get_logger('test_eval') + runner = EpochBasedRunner( + model=model, + batch_processor=None, + optimizer=optimizer, + work_dir=tmpdir, + logger=logger) + runner.register_checkpoint_hook(dict(interval=1)) + runner.register_hook(eval_hook) + runner.run([loader], [('train', 1)], 8) + + real_path = osp.join(tmpdir, 'best_mpjpe_epoch_4.pth') + + assert runner.meta['hook_msgs']['best_ckpt'] == osp.realpath(real_path) + assert runner.meta['hook_msgs']['best_score'] == 0.05 + + data_loader = DataLoader(EvalDataset(), batch_size=1) + eval_hook = EvalHookCls( + data_loader, interval=1, save_best='pve', rule='less') + with tempfile.TemporaryDirectory() as tmpdir: + logger = get_logger('test_eval') + runner = EpochBasedRunner( + model=model, + batch_processor=None, + optimizer=optimizer, + work_dir=tmpdir, + logger=logger) + runner.register_checkpoint_hook(dict(interval=1)) + runner.register_hook(eval_hook) + runner.run([loader], [('train', 1)], 8) + + real_path = osp.join(tmpdir, 'best_pve_epoch_4.pth') + + assert runner.meta['hook_msgs']['best_ckpt'] == osp.realpath(real_path) + assert runner.meta['hook_msgs']['best_score'] == 0.05 + + data_loader = DataLoader(EvalDataset(), batch_size=1) + eval_hook = EvalHookCls(data_loader, save_best='mpjpe', rule='greater') + with tempfile.TemporaryDirectory() as tmpdir: + logger = get_logger('test_eval') + runner = EpochBasedRunner( + model=model, + batch_processor=None, + optimizer=optimizer, + work_dir=tmpdir, + logger=logger, + max_epochs=8) + runner.register_checkpoint_hook(dict(interval=1)) + runner.register_hook(eval_hook) + runner.run([loader], [('train', 1)]) + + real_path = osp.join(tmpdir, 'best_mpjpe_epoch_6.pth') + + assert runner.meta['hook_msgs']['best_ckpt'] == osp.realpath(real_path) + assert runner.meta['hook_msgs']['best_score'] == 0.7 + + data_loader = DataLoader(EvalDataset(), batch_size=1) + eval_hook = EvalHookCls(data_loader, save_best='mpjpe') + with tempfile.TemporaryDirectory() as tmpdir: + logger = get_logger('test_eval') + runner = EpochBasedRunner( + model=model, + batch_processor=None, + optimizer=optimizer, + work_dir=tmpdir, + logger=logger, + max_epochs=2) + runner.register_checkpoint_hook(dict(interval=1)) + runner.register_hook(eval_hook) + runner.run([loader], [('train', 1)]) + + real_path = osp.join(tmpdir, 'best_mpjpe_epoch_1.pth') + + assert runner.meta['hook_msgs']['best_ckpt'] == osp.realpath(real_path) + assert runner.meta['hook_msgs']['best_score'] == 0.1 + + resume_from = osp.join(tmpdir, 'latest.pth') + loader = DataLoader(ExampleDataset(), batch_size=1) + eval_hook = EvalHookCls(data_loader, save_best='mpjpe') + runner = EpochBasedRunner( + model=model, + batch_processor=None, + optimizer=optimizer, + work_dir=tmpdir, + logger=logger, + max_epochs=8) + runner.register_checkpoint_hook(dict(interval=1)) + runner.register_hook(eval_hook) + runner.resume(resume_from) + runner.run([loader], [('train', 1)]) + + real_path = osp.join(tmpdir, 'best_mpjpe_epoch_4.pth') + + assert runner.meta['hook_msgs']['best_ckpt'] == osp.realpath(real_path) + assert runner.meta['hook_msgs']['best_score'] == 0.05 diff --git a/tests/test_evaluation/test_eval_utils.py b/tests/test_evaluation/test_eval_utils.py new file mode 100644 index 00000000..6f0dfbf6 --- /dev/null +++ b/tests/test_evaluation/test_eval_utils.py @@ -0,0 +1,101 @@ +import numpy as np +import pytest + +from mmhuman3d.core.evaluation import ( + keypoint_3d_auc, + keypoint_3d_pck, + keypoint_accel_error, + keypoint_mpjpe, + vertice_pve, +) + + +def tets_accel_error(): + target = np.random.rand(10, 5, 3) + output = np.copy(target) + mask = np.ones((output.shape[0]), dtype=bool) + + error = keypoint_accel_error(output, target, mask) + np.testing.assert_almost_equal(error, 0) + + error = keypoint_accel_error(output, target) + np.testing.assert_almost_equal(error, 0) + + +def tets_keypoinyt_mpjpe(): + target = np.random.rand(2, 5, 3) + output = np.copy(target) + mask = np.ones((output.shape[0], output.shape[1]), dtype=bool) + with pytest.raises(ValueError): + _ = keypoint_mpjpe(output, target, mask, alignment='norm') + + error = keypoint_mpjpe(output, target, mask, alignment='none') + np.testing.assert_almost_equal(error, 0) + + error = keypoint_mpjpe(output, target, mask, alignment='scale') + np.testing.assert_almost_equal(error, 0) + + error = keypoint_mpjpe(output, target, mask, alignment='procrustes') + np.testing.assert_almost_equal(error, 0) + + R = np.array([[1, 0, 0], [0, -1, 0], [0, 0, -1]]) + output = np.dot(target, R) + error = keypoint_mpjpe(output, target, mask, alignment='none') + assert error > 1e-10 + + R = np.array([[1, 0, 0], [0, -1, 0], [0, 0, -1]]) + output = np.dot(target, R) + error = keypoint_mpjpe(output, target, mask, alignment='procrustes') + np.testing.assert_almost_equal(error, 0) + + +def tets_keypoinyt_pve(): + target = np.random.rand(2, 6890, 3) + output = np.copy(target) + + error = vertice_pve(output, target) + np.testing.assert_almost_equal(error, 0) + + +def test_keypoint_3d_pck(): + target = np.random.rand(2, 5, 3) * 1000 + output = np.copy(target) + mask = np.ones((output.shape[0], output.shape[1]), dtype=bool) + + with pytest.raises(ValueError): + _ = keypoint_3d_pck(output, target, mask, alignment='norm') + + pck = keypoint_3d_pck(output, target, mask, alignment='none') + np.testing.assert_almost_equal(pck, 100) + + output[0, 0, :] = target[0, 0, :] + 1000 + pck = keypoint_3d_pck(output, target, mask, alignment='none') + np.testing.assert_almost_equal(pck, 90, 5) + + output = target * 2 + pck = keypoint_3d_pck(output, target, mask, alignment='scale') + np.testing.assert_almost_equal(pck, 100) + + output = target + 2 + pck = keypoint_3d_pck(output, target, mask, alignment='procrustes') + np.testing.assert_almost_equal(pck, 100) + + +def test_keypoint_3d_auc(): + target = np.random.rand(2, 5, 3) * 1000 + output = np.copy(target) + mask = np.ones((output.shape[0], output.shape[1]), dtype=bool) + + with pytest.raises(ValueError): + _ = keypoint_3d_auc(output, target, mask, alignment='norm') + + auc = keypoint_3d_auc(output, target, mask, alignment='none') + np.testing.assert_almost_equal(auc, 30 / 31 * 100) + + output = target * 2 + auc = keypoint_3d_auc(output, target, mask, alignment='scale') + np.testing.assert_almost_equal(auc, 30 / 31 * 100) + + output = target + 2000 + auc = keypoint_3d_auc(output, target, mask, alignment='procrustes') + np.testing.assert_almost_equal(auc, 30 / 31 * 100) diff --git a/tests/test_hybrik_pipeline.py b/tests/test_hybrik_pipeline.py index 4cbe68da..c0adac19 100644 --- a/tests/test_hybrik_pipeline.py +++ b/tests/test_hybrik_pipeline.py @@ -397,10 +397,12 @@ def test_human_hybrik_dataset(): dataset = 'HybrIKHumanImageDataset' dataset_class = DATASETS.get(dataset) + body_model = dict(type='SMPL', model_path='data/body_models/smpl') # train mode custom_dataset = dataset_class( dataset_name='h36m', data_prefix='tests/data', + body_model=body_model, pipeline=[], ann_file='h36m_hybrik_train.npz') @@ -415,23 +417,50 @@ def test_human_hybrik_dataset(): sample_item = custom_dataset[0] for k in keys: assert k in sample_item - + body_model = dict(type='SMPL', model_path='data/body_models/smpl') # test mode + num_data = 1 custom_dataset = dataset_class( dataset_name='h36m', data_prefix='tests/data', + body_model=body_model, pipeline=[], ann_file='h36m_hybrik_train.npz', test_mode=True) - + custom_dataset.num_data = num_data # test evaluation - outputs = [] - for item in custom_dataset: - pred = dict( - xyz_17=item['joint_relative_17'][None, ...], - image_idx=[item['sample_idx']]) - outputs.append(pred) + outputs = [{ + 'xyz_17': np.random.rand(num_data, 17, 3), + 'smpl_pose': np.random.rand(num_data, 24, 3, 3), + 'smpl_beta': np.random.rand(num_data, 10), + 'image_idx': np.arange(num_data) + }] with tempfile.TemporaryDirectory() as tmpdir: eval_result = custom_dataset.evaluate(outputs, tmpdir) - assert 'MPJPE' in eval_result - assert 'MPJPE-PA' in eval_result + assert 'PA-MPJPE' in eval_result + assert eval_result['PA-MPJPE'] > 0 + + res = custom_dataset.evaluate( + outputs, res_folder=tmpdir, metric='mpjpe') + assert 'MPJPE' in res + assert res['MPJPE'] > 0 + + res = custom_dataset.evaluate( + outputs, res_folder=tmpdir, metric='pa-3dpck') + assert 'PA-3DPCK' in res + assert res['PA-3DPCK'] >= 0 + + res = custom_dataset.evaluate( + outputs, res_folder=tmpdir, metric='3dpck') + assert '3DPCK' in res + assert res['3DPCK'] >= 0 + + res = custom_dataset.evaluate( + outputs, res_folder=tmpdir, metric='pa-3dauc') + assert 'PA-3DAUC' in res + assert res['PA-3DAUC'] >= 0 + + res = custom_dataset.evaluate( + outputs, res_folder=tmpdir, metric='3dauc') + assert '3DAUC' in res + assert res['3DAUC'] >= 0 diff --git a/tools/test.py b/tools/test.py index 3c8aeda1..a93cdedf 100644 --- a/tools/test.py +++ b/tools/test.py @@ -1,5 +1,6 @@ import argparse import os +import os.path as osp import mmcv import torch @@ -28,38 +29,27 @@ def parse_args(): '--metrics', type=str, nargs='+', - help='evaluation metrics, which depends on the dataset, e.g., ' - '"accuracy", "precision", "recall", "f1_score", "support" for single ' - 'label dataset, and "mAP", "CP", "CR", "CF1", "OP", "OR", "OF1" for ' - 'multi-label dataset') - parser.add_argument('--show', action='store_true', help='show results') - parser.add_argument( - '--show-dir', help='directory where painted images will be saved') + default='pa-mpjpe', + help='evaluation metric, which depends on the dataset,' + ' e.g., "pa-mpjpe" for H36M') parser.add_argument( '--gpu_collect', action='store_true', help='whether to use gpu to collect results') parser.add_argument('--tmpdir', help='tmp dir for writing some results') parser.add_argument( - '--options', + '--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.') parser.add_argument( - '--metric-options', + '--eval-options', nargs='+', - action=DictAction, default={}, - help='custom options for evaluation, the key-value pair in xxx=yyy ' - 'format will be parsed as a dict metric_options for dataset.evaluate()' - ' function.') - parser.add_argument( - '--show-options', - nargs='+', action=DictAction, - help='custom options for show_result. key-value pair in xxx=yyy.' - 'Check available options in `model.show_result`.') + help='custom options for evaluation, the key-value pair in xxx=yyy ' + 'format will be kwargs for dataset.evaluate() function') parser.add_argument( '--launcher', choices=['none', 'pytorch', 'slurm', 'mpi'], @@ -81,8 +71,8 @@ def main(): args = parse_args() cfg = mmcv.Config.fromfile(args.config) - if args.options is not None: - cfg.merge_from_dict(args.options) + if args.cfg_options is not None: + cfg.merge_from_dict(args.cfg_options) # set cudnn_benchmark if cfg.get('cudnn_benchmark', False): torch.backends.cudnn.benchmark = True @@ -118,9 +108,7 @@ def main(): model = model.cpu() else: model = MMDataParallel(model, device_ids=[0]) - show_kwargs = {} if args.show_options is None else args.show_options - outputs = single_gpu_test(model, data_loader, args.show, args.show_dir, - **show_kwargs) + outputs = single_gpu_test(model, data_loader) else: model = MMDistributedDataParallel( model.cuda(), @@ -130,8 +118,10 @@ def main(): args.gpu_collect) rank, _ = get_dist_info() - eval_cfg = cfg.get('evaluation', {}) + eval_cfg = cfg.get('evaluation', args.eval_options) + eval_cfg.update(dict(metric=args.metrics)) if rank == 0: + mmcv.mkdir_or_exist(osp.abspath(args.work_dir)) results = dataset.evaluate(outputs, args.work_dir, **eval_cfg) for k, v in results.items(): print(f'\n{k} : {v:.2f}')