diff --git a/mmpose/datasets/transforms/converting.py b/mmpose/datasets/transforms/converting.py index fc8b335e26..bca000435a 100644 --- a/mmpose/datasets/transforms/converting.py +++ b/mmpose/datasets/transforms/converting.py @@ -144,44 +144,6 @@ def transform(self, results: dict) -> dict: return results - def transform_sigmas(self, sigmas: Union[List, np.ndarray]): - """Transforms the sigmas based on the mapping.""" - list_input = False - if isinstance(sigmas, list): - sigmas = np.array(sigmas) - list_input = True - - new_sigmas = np.ones(self.num_keypoints, dtype=sigmas.dtype) - new_sigmas[self.target_index] = sigmas[self.source_index] - - if list_input: - new_sigmas = new_sigmas.tolist() - - return new_sigmas - - def transform_ann(self, ann_info: Union[dict, list]): - """Transforms the annotations based on the mapping.""" - - list_input = True - if not isinstance(ann_info, list): - ann_info = [ann_info] - list_input = False - - for ann in ann_info: - if 'keypoints' in ann: - keypoints = np.array(ann['keypoints']).reshape(-1, 3) - new_keypoints = np.zeros((self.num_keypoints, 3), - dtype=keypoints.dtype) - new_keypoints[self.target_index] = keypoints[self.source_index] - ann['keypoints'] = new_keypoints.reshape(-1).tolist() - if 'num_keypoints' in ann: - ann['num_keypoints'] = self.num_keypoints - - if not list_input: - ann_info = ann_info[0] - - return ann_info - def __repr__(self) -> str: """print the basic information of the transform. diff --git a/mmpose/evaluation/functional/__init__.py b/mmpose/evaluation/functional/__init__.py index f59f6f5a40..47255fc394 100644 --- a/mmpose/evaluation/functional/__init__.py +++ b/mmpose/evaluation/functional/__init__.py @@ -4,10 +4,11 @@ multilabel_classification_accuracy, pose_pck_accuracy, simcc_pck_accuracy) from .nms import nms, nms_torch, oks_nms, soft_oks_nms +from .transforms import transform_ann, transform_pred, transform_sigmas __all__ = [ 'keypoint_pck_accuracy', 'keypoint_auc', 'keypoint_nme', 'keypoint_epe', 'pose_pck_accuracy', 'multilabel_classification_accuracy', 'simcc_pck_accuracy', 'nms', 'oks_nms', 'soft_oks_nms', 'keypoint_mpjpe', - 'nms_torch' + 'nms_torch', 'transform_ann', 'transform_sigmas', 'transform_pred' ] diff --git a/mmpose/evaluation/functional/transforms.py b/mmpose/evaluation/functional/transforms.py new file mode 100644 index 0000000000..56873b389c --- /dev/null +++ b/mmpose/evaluation/functional/transforms.py @@ -0,0 +1,99 @@ +# Copyright (c) OpenMMLab. All rights reserved. +from typing import List, Tuple, Union + +import numpy as np + + +def transform_sigmas(sigmas: Union[List, np.ndarray], num_keypoints: int, + mapping: Union[List[Tuple[int, int]], List[Tuple[Tuple, + int]]]): + """Transforms the sigmas based on the mapping.""" + if len(mapping): + source_index, target_index = map(list, zip(*mapping)) + else: + source_index, target_index = [], [] + + list_input = False + if isinstance(sigmas, list): + sigmas = np.array(sigmas) + list_input = True + + new_sigmas = np.ones(num_keypoints, dtype=sigmas.dtype) + new_sigmas[target_index] = sigmas[source_index] + + if list_input: + new_sigmas = new_sigmas.tolist() + + return new_sigmas + + +def transform_ann(ann_info: Union[dict, list], num_keypoints: int, + mapping: Union[List[Tuple[int, int]], List[Tuple[Tuple, + int]]]): + """Transforms COCO-format annotations based on the mapping.""" + if len(mapping): + source_index, target_index = map(list, zip(*mapping)) + else: + source_index, target_index = [], [] + + list_input = True + if not isinstance(ann_info, list): + ann_info = [ann_info] + list_input = False + + for each in ann_info: + if 'keypoints' in each: + keypoints = np.array(each['keypoints']) + + C = 3 # COCO-format: x, y, score + keypoints = keypoints.reshape(-1, C) + new_keypoints = np.zeros((num_keypoints, C), dtype=keypoints.dtype) + new_keypoints[target_index] = keypoints[source_index] + each['keypoints'] = new_keypoints.reshape(-1).tolist() + + if 'num_keypoints' in each: + each['num_keypoints'] = num_keypoints + + if not list_input: + ann_info = ann_info[0] + + return ann_info + + +def transform_pred(pred_info: Union[dict, list], num_keypoints: int, + mapping: Union[List[Tuple[int, int]], List[Tuple[Tuple, + int]]]): + """Transforms predictions based on the mapping.""" + if len(mapping): + source_index, target_index = map(list, zip(*mapping)) + else: + source_index, target_index = [], [] + + list_input = True + if not isinstance(pred_info, list): + pred_info = [pred_info] + list_input = False + + for each in pred_info: + if 'keypoints' in each: + keypoints = np.array(each['keypoints']) + + N, _, C = keypoints.shape + new_keypoints = np.zeros((N, num_keypoints, C), + dtype=keypoints.dtype) + new_keypoints[:, target_index] = keypoints[:, source_index] + each['keypoints'] = new_keypoints + + keypoint_scores = np.array(each['keypoint_scores']) + new_scores = np.zeros((N, num_keypoints), + dtype=keypoint_scores.dtype) + new_scores[:, target_index] = keypoint_scores[:, source_index] + each['keypoint_scores'] = new_scores + + if 'num_keypoints' in each: + each['num_keypoints'] = num_keypoints + + if not list_input: + pred_info = pred_info[0] + + return pred_info diff --git a/mmpose/evaluation/metrics/coco_metric.py b/mmpose/evaluation/metrics/coco_metric.py index 8fc32dd809..84528041e7 100644 --- a/mmpose/evaluation/metrics/coco_metric.py +++ b/mmpose/evaluation/metrics/coco_metric.py @@ -12,9 +12,10 @@ from xtcocotools.coco import COCO from xtcocotools.cocoeval import COCOeval -from mmpose.registry import METRICS, TRANSFORMS +from mmpose.registry import METRICS from mmpose.structures.bbox import bbox_xyxy2xywh -from ..functional import oks_nms, soft_oks_nms +from ..functional import (oks_nms, soft_oks_nms, transform_ann, transform_pred, + transform_sigmas) @METRICS.register_module() @@ -73,10 +74,12 @@ class CocoMetric(BaseMetric): test submission when the ground truth annotations are absent. If set to ``True``, ``outfile_prefix`` should specify the path to store the output results. Defaults to ``False`` + pred_converter (dict, optional): Config dictionary for the prediction + converter. The dictionary has the same parameters as + 'KeypointConverter'. Defaults to None. gt_converter (dict, optional): Config dictionary for the ground truth - converter. The dictionary must contain the key 'type' set to - 'KeypointConverter' to indicate the type of ground truth converter - to be used. Defaults to None. + converter. The dictionary has the same parameters as + 'KeypointConverter'. Defaults to None. outfile_prefix (str | None): The prefix of json files. It includes the file path and the prefix of filename, e.g., ``'a/b/prefix'``. If not specified, a temp file will be created. Defaults to ``None`` @@ -99,7 +102,8 @@ def __init__(self, nms_mode: str = 'oks_nms', nms_thr: float = 0.9, format_only: bool = False, - gt_converter: Optional[dict] = None, + pred_converter: Dict = None, + gt_converter: Dict = None, outfile_prefix: Optional[str] = None, collect_device: str = 'cpu', prefix: Optional[str] = None) -> None: @@ -145,13 +149,8 @@ def __init__(self, self.format_only = format_only self.outfile_prefix = outfile_prefix - - if gt_converter is not None: - assert gt_converter.get('type', None) == 'KeypointConverter', \ - 'the type of `gt_converter` must be \'KeypointConverter\'' - self.gt_converter = TRANSFORMS.build(gt_converter) - else: - self.gt_converter = None + self.pred_converter = pred_converter + self.gt_converter = gt_converter @property def dataset_meta(self) -> Optional[dict]: @@ -162,8 +161,9 @@ def dataset_meta(self) -> Optional[dict]: def dataset_meta(self, dataset_meta: dict) -> None: """Set the dataset meta info to the metric.""" if self.gt_converter is not None: - dataset_meta['sigmas'] = self.gt_converter.transform_sigmas( - dataset_meta['sigmas']) + dataset_meta['sigmas'] = transform_sigmas( + dataset_meta['sigmas'], self.gt_converter['num_keypoints'], + self.gt_converter['mapping']) dataset_meta['num_keypoints'] = len(dataset_meta['sigmas']) self._dataset_meta = dataset_meta @@ -394,19 +394,28 @@ def compute_metrics(self, results: list) -> Dict[str, float]: self.coco = COCO(coco_json_path) if self.gt_converter is not None: for id_, ann in self.coco.anns.items(): - self.coco.anns[id_] = self.gt_converter.transform_ann(ann) + self.coco.anns[id_] = transform_ann( + ann, self.gt_converter['num_keypoints'], + self.gt_converter['mapping']) kpts = defaultdict(list) # group the preds by img_id for pred in preds: img_id = pred['img_id'] - for idx in range(len(pred['keypoints'])): + + if self.pred_converter is not None: + pred = transform_pred(pred, + self.pred_converter['num_keypoints'], + self.pred_converter['mapping']) + + for idx, keypoints in enumerate(pred['keypoints']): + instance = { 'id': pred['id'], 'img_id': pred['img_id'], 'category_id': pred['category_id'], - 'keypoints': pred['keypoints'][idx], + 'keypoints': keypoints, 'keypoint_scores': pred['keypoint_scores'][idx], 'bbox_score': pred['bbox_scores'][idx], } @@ -417,7 +426,6 @@ def compute_metrics(self, results: list) -> Dict[str, float]: instance['area'] = pred['areas'][idx] else: # use keypoint to calculate bbox and get area - keypoints = pred['keypoints'][idx] area = ( np.max(keypoints[:, 0]) - np.min(keypoints[:, 0])) * ( np.max(keypoints[:, 1]) - np.min(keypoints[:, 1])) @@ -431,7 +439,10 @@ def compute_metrics(self, results: list) -> Dict[str, float]: # score the prediction results according to `score_mode` # and perform NMS according to `nms_mode` valid_kpts = defaultdict(list) - num_keypoints = self.dataset_meta['num_keypoints'] + if self.pred_converter is not None: + num_keypoints = self.pred_converter['num_keypoints'] + else: + num_keypoints = self.dataset_meta['num_keypoints'] for img_id, instances in kpts.items(): for instance in instances: # concatenate the keypoint coordinates and scores diff --git a/tests/test_datasets/test_transforms/test_converting.py b/tests/test_datasets/test_transforms/test_converting.py index dc4376baf9..9f535b5126 100644 --- a/tests/test_datasets/test_transforms/test_converting.py +++ b/tests/test_datasets/test_transforms/test_converting.py @@ -1,5 +1,5 @@ # Copyright (c) OpenMMLab. All rights reserved. -from copy import deepcopy + from unittest import TestCase import numpy as np @@ -107,30 +107,3 @@ def test_transform(self): self.assertTrue( (results['keypoints_visible'][:, target_index, 0] == self.data_info['keypoints_visible'][:, source_index]).all()) - - def test_transform_sigmas(self): - - mapping = [(3, 0), (6, 1), (16, 2), (5, 3)] - transform = KeypointConverter(num_keypoints=5, mapping=mapping) - sigmas = np.random.rand(17) - new_sigmas = transform.transform_sigmas(sigmas) - self.assertEqual(len(new_sigmas), 5) - for i, j in mapping: - self.assertEqual(sigmas[i], new_sigmas[j]) - - def test_transform_ann(self): - mapping = [(3, 0), (6, 1), (16, 2), (5, 3)] - transform = KeypointConverter(num_keypoints=5, mapping=mapping) - - ann_info = dict( - num_keypoints=17, - keypoints=np.random.randint(3, size=(17 * 3, )).tolist()) - ann_info_copy = deepcopy(ann_info) - - _ = transform.transform_ann(ann_info) - - self.assertEqual(ann_info['num_keypoints'], 5) - self.assertEqual(len(ann_info['keypoints']), 15) - for i, j in mapping: - self.assertListEqual(ann_info_copy['keypoints'][i * 3:i * 3 + 3], - ann_info['keypoints'][j * 3:j * 3 + 3]) diff --git a/tests/test_evaluation/test_functional/test_transforms.py b/tests/test_evaluation/test_functional/test_transforms.py new file mode 100644 index 0000000000..22ad39b604 --- /dev/null +++ b/tests/test_evaluation/test_functional/test_transforms.py @@ -0,0 +1,56 @@ +# Copyright (c) OpenMMLab. All rights reserved. +from copy import deepcopy +from unittest import TestCase + +import numpy as np + +from mmpose.evaluation.functional import (transform_ann, transform_pred, + transform_sigmas) + + +class TestKeypointEval(TestCase): + + def test_transform_sigmas(self): + + mapping = [(3, 0), (6, 1), (16, 2), (5, 3)] + num_keypoints = 5 + sigmas = np.random.rand(17) + new_sigmas = transform_sigmas(sigmas, num_keypoints, mapping) + self.assertEqual(len(new_sigmas), 5) + for i, j in mapping: + self.assertEqual(sigmas[i], new_sigmas[j]) + + def test_transform_ann(self): + mapping = [(3, 0), (6, 1), (16, 2), (5, 3)] + num_keypoints = 5 + + kpt_info = dict( + num_keypoints=17, + keypoints=np.random.randint(3, size=(17 * 3, )).tolist()) + kpt_info_copy = deepcopy(kpt_info) + + _ = transform_ann(kpt_info, num_keypoints, mapping) + + self.assertEqual(kpt_info['num_keypoints'], 5) + self.assertEqual(len(kpt_info['keypoints']), 15) + for i, j in mapping: + self.assertListEqual(kpt_info_copy['keypoints'][i * 3:i * 3 + 3], + kpt_info['keypoints'][j * 3:j * 3 + 3]) + + def test_transform_pred(self): + mapping = [(3, 0), (6, 1), (16, 2), (5, 3)] + num_keypoints = 5 + + kpt_info = dict( + num_keypoints=17, + keypoints=np.random.randint(3, size=( + 1, + 17, + 3, + )), + keypoint_scores=np.ones((1, 17))) + + _ = transform_pred(kpt_info, num_keypoints, mapping) + + self.assertEqual(kpt_info['num_keypoints'], 5) + self.assertEqual(len(kpt_info['keypoints']), 1)