fix wholebody keypoints eval

easycv/core/evaluation/wholebody_keypoint_eval.py

@@ -1,34 +1,45 @@
 # Copyright (c) OpenMMLab. All rights reserved.
 # Adapt from
 # https://github.com/open-mmlab/mmpose/blob/master/mmpose/datasets/datasets/base/kpt_2d_sview_rgb_img_top_down_dataset.py
+import json
+import os
+import tempfile
+
 import numpy as np
+from xtcocotools.coco import COCO
 from xtcocotools.cocoeval import COCOeval
 
+from easycv.utils.json_utils import MyEncoder
 from .builder import EVALUATORS
 from .coco_evaluation import CoCoPoseTopDownEvaluator
 from .metric_registry import METRICS
-from .top_down_eval import (keypoint_auc, keypoint_epe, keypoint_nme,
-                            keypoint_pck_accuracy)
 
 
 @EVALUATORS.register_module
 class WholeBodyKeyPointEvaluator(CoCoPoseTopDownEvaluator):
     """ KeyPoint evaluator.
     """
 
-    def __init__(self, dataset_name=None, metric_names=['AP']):
+    def __init__(self, dataset_name=None, metric_names=['AP'], **kwargs):
         """
 
         Args:
             dataset_name: eval dataset name
             metric_names: eval metrics name
         """
         super(WholeBodyKeyPointEvaluator,
-              self).__init__(dataset_name, metric_names)
+              self).__init__(dataset_name, metric_names, **kwargs)
         self.metric = metric_names
         self.dataset_name = dataset_name
-
-    def _coco_keypoint_results_one_category_kernel(self, data_pack):
+        self.body_num = kwargs.get('body_num', 17)
+        self.foot_num = kwargs.get('foot_num', 6)
+        self.face_num = kwargs.get('face_num', 68)
+        self.left_hand_num = kwargs.get('left_hand_num', 21)
+        self.right_hand_num = kwargs.get('right_hand_num', 21)
+
+    def _coco_keypoint_results_one_category_kernel(self,
+                                                   data_pack,
+                                                   num_joints=None):
         """Get coco keypoint results."""
         cat_id = data_pack['cat_id']
         keypoints = data_pack['keypoints']
@@ -40,8 +51,7 @@ def _coco_keypoint_results_one_category_kernel(self, data_pack):
 
             _key_points = np.array(
                 [img_kpt['keypoints'] for img_kpt in img_kpts])
-            key_points = _key_points.reshape(-1,
-                                             self.ann_info['num_joints'] * 3)
+            key_points = _key_points.reshape(-1, num_joints * 3)
 
             cuts = np.cumsum([
                 0, self.body_num, self.foot_num, self.face_num,
@@ -65,76 +75,82 @@ def _coco_keypoint_results_one_category_kernel(self, data_pack):
 
         return cat_results
 
-    def _do_python_keypoint_eval(self, res_file):
+    def _do_python_keypoint_eval(self, results, groundtruth, sigmas=None):
         """Keypoint evaluation using COCOAPI."""
-        coco_det = self.coco.loadRes(res_file)
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            groundtruth_file = os.path.join(
+                tmp_dir, 'groundtruth_wholebody_keypoints.json')
+            with open(groundtruth_file, 'w') as f:
+                json.dump(groundtruth, f, sort_keys=True, indent=4)
+            coco = COCO(groundtruth_file)
+
+            res_file = os.path.join(tmp_dir, 'result_wholebody_keypoints.json')
+            with open(res_file, 'w') as f:
+                json.dump(results, f, sort_keys=True, indent=4, cls=MyEncoder)
+            coco_det = coco.loadRes(res_file)
 
         cuts = np.cumsum([
             0, self.body_num, self.foot_num, self.face_num, self.left_hand_num,
             self.right_hand_num
         ])
 
         coco_eval = COCOeval(
-            self.coco,
+            coco,
             coco_det,
             'keypoints_body',
-            self.sigmas[cuts[0]:cuts[1]],
+            sigmas[cuts[0]:cuts[1]],
             use_area=True)
         coco_eval.params.useSegm = None
         coco_eval.evaluate()
         coco_eval.accumulate()
         coco_eval.summarize()
 
         coco_eval = COCOeval(
-            self.coco,
+            coco,
             coco_det,
             'keypoints_foot',
-            self.sigmas[cuts[1]:cuts[2]],
+            sigmas[cuts[1]:cuts[2]],
             use_area=True)
         coco_eval.params.useSegm = None
         coco_eval.evaluate()
         coco_eval.accumulate()
         coco_eval.summarize()
 
         coco_eval = COCOeval(
-            self.coco,
+            coco,
             coco_det,
             'keypoints_face',
-            self.sigmas[cuts[2]:cuts[3]],
+            sigmas[cuts[2]:cuts[3]],
             use_area=True)
         coco_eval.params.useSegm = None
         coco_eval.evaluate()
         coco_eval.accumulate()
         coco_eval.summarize()
 
         coco_eval = COCOeval(
-            self.coco,
+            coco,
             coco_det,
             'keypoints_lefthand',
-            self.sigmas[cuts[3]:cuts[4]],
+            sigmas[cuts[3]:cuts[4]],
             use_area=True)
         coco_eval.params.useSegm = None
         coco_eval.evaluate()
         coco_eval.accumulate()
         coco_eval.summarize()
 
         coco_eval = COCOeval(
-            self.coco,
+            coco,
             coco_det,
             'keypoints_righthand',
-            self.sigmas[cuts[4]:cuts[5]],
+            sigmas[cuts[4]:cuts[5]],
             use_area=True)
         coco_eval.params.useSegm = None
         coco_eval.evaluate()
         coco_eval.accumulate()
         coco_eval.summarize()
 
         coco_eval = COCOeval(
-            self.coco,
-            coco_det,
-            'keypoints_wholebody',
-            self.sigmas,
-            use_area=True)
+            coco, coco_det, 'keypoints_wholebody', sigmas, use_area=True)
         coco_eval.params.useSegm = None
         coco_eval.evaluate()
         coco_eval.accumulate()

easycv/datasets/pose/wholebody_topdown_coco_dataset.py

@@ -28,26 +28,20 @@ def evaluate(self, outputs, evaluators, **kwargs):
         if len(evaluators) > 1 or not isinstance(evaluators[0],
                                                  WholeBodyKeyPointEvaluator):
             raise ValueError(
-                'HandCocoWholeBodyDataset only support one `WholeBodyKeyPointEvaluator` now, '
+                'WholeBodyCocoTopDownDataset only support one `WholeBodyKeyPointEvaluator` now, '
                 'but get %s' % evaluators)
-        evaluator = evaluators[0]
-
-        image_ids = outputs['image_ids']
-        preds = outputs['preds']
-        boxes = outputs['boxes']
-        bbox_ids = outputs['bbox_ids']
-
-        kpts = []
-        for i, image_id in enumerate(image_ids):
-            kpts.append({
-                'keypoints': preds[i],
-                'center': boxes[i][0:2],
-                'scale': boxes[i][2:4],
-                'area': boxes[i][4],
-                'score': boxes[i][5],
-                'image_id': image_id,
-                'bbox_id': bbox_ids[i]
-            })
-        kpts = self._sort_and_unique_bboxes(kpts)
-        eval_res = evaluator.evaluate(kpts, self.data_source.db)
-        return eval_res
+
+        evaluator_args = {
+            'num_joints': self.data_source.ann_info['num_joints'],
+            'sigmas': self.data_source.sigmas,
+            'class2id': self.data_source._class_to_ind
+        }
+        eval_result = {}
+        for evaluator in evaluators:
+            eval_result.update(
+                evaluator.evaluate(
+                    prediction_dict=outputs,
+                    groundtruth_dict=self.data_source.coco.dataset,
+                    **evaluator_args))
+
+        return eval_result