text: # Copyright (c) 2019-present, Baidu, Inc.
  #
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #
  #     http://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  ##############################################################################
  
  """Interface for COCO evaluation."""
  
  import os
  import json
  
  import numpy as np
  from collections import defaultdict
  from collections import OrderedDict
  import pickle
  
  from utils.base_evaluator import BaseEvaluator
  from pycocotools.coco import COCO
  from pycocotools.cocoeval import COCOeval
  from nms.nms import oks_nms
  
  
  class COCOEvaluator(BaseEvaluator):
      def __init__(self, root, kp_dim=17):
          """
          :param root: the root dir of dataset
          :param kp_dim: the dimension of keypoints
          """
          super(COCOEvaluator, self).__init__(root, kp_dim)
          self.kp_dim = kp_dim
          self.in_vis_thre = 0.2
          self.oks_thre = 0.9
          self.coco = COCO(os.path.join(root, 'annotations', 'person_keypoints_val2017.json'))
  
          cats = [cat['name']
                  for cat in self.coco.loadCats(self.coco.getCatIds())]
          self.classes = ['__background__'] + cats
          self.num_classes = len(self.classes)
          self._class_to_ind = dict(zip(self.classes, range(self.num_classes)))
          self._class_to_coco_ind = dict(zip(cats, self.coco.getCatIds()))
          self._coco_ind_to_class_ind = dict([(self._class_to_coco_ind[cls],
                                               self._class_to_ind[cls])
                                              for cls in self.classes[1:]])
  
      def evaluate(self, preds, output_dir, all_boxes, img_path, *args, **kwargs):
          """
          :param preds: the predictions to be evaluated
          :param output_dir: target directory to save evaluation results
          :param all_boxes: ground truth
          :param img_path: paths of the original image
          :return:
          """
          res_folder = os.path.join(output_dir, 'results')
          if not os.path.exists(res_folder):
              os.makedirs(res_folder)
          res_file = os.path.join(res_folder, 'keypoints_coco_results.json')
  
          # person x (keypoints)
          _kpts = []
          for idx, kpt in enumerate(preds):
              _kpts.append({
                  'keypoints': kpt,
                  'center': all_boxes[idx][0:2],
                  'scale': all_boxes[idx][2:4],
                  'area': all_boxes[idx][4],
                  'score': all_boxes[idx][5],
                  'image': int(img_path[idx][-16:-4])
              })
          # image x person x (keypoints)
          kpts = defaultdict(list)
          for kpt in _kpts:
              kpts[kpt['image']].append(kpt)
  
          # rescoring and oks nms
          kp_dim = self.kp_dim
          in_vis_thre = self.in_vis_thre
          oks_thre = self.oks_thre
          oks_nmsed_kpts = []
          for img in kpts.keys():
              img_kpts = kpts[img]
              for n_p in img_kpts:
                  box_score = n_p['score']
                  kpt_score = 0
                  valid_num = 0
                  for n_jt in range(0, kp_dim):
                      t_s = n_p['keypoints'][n_jt][2]
                      if t_s > in_vis_thre:
                          kpt_score = kpt_score + t_s
                          valid_num = valid_num + 1
                  if valid_num != 0:
                      kpt_score = kpt_score / valid_num
                  # rescoring
                  n_p['score'] = kpt_score * box_score
              keep = oks_nms([img_kpts[i] for i in range(len(img_kpts))], oks_thre)
              if len(keep) == 0:
                  oks_nmsed_kpts.append(img_kpts)
              else:
                  oks_nmsed_kpts.append([img_kpts[_keep] for _keep in keep])
  
          self._write_coco_keypoint_results(oks_nmsed_kpts, res_file)
          info_str = self._do_python_keypoint_eval(res_file, res_folder)
          name_value = OrderedDict(info_str)
          return name_value, name_value['AP']
  
      def _write_coco_keypoint_results(self, keypoints, res_file):
          data_pack = [{'cat_id': self._class_to_coco_ind[cls],
                        'cls_ind': cls_ind,
                        'cls': cls,
                        'ann_type': 'keypoints',
                        'keypoints': keypoints
                        }
                       for cls_ind, cls in enumerate(self.classes) if not cls == '__background__']
  
          results = self._coco_keypoint_results_one_category_kernel(data_pack[0])
          with open(res_file, 'w') as f:
              json.dump(results, f, sort_keys=True, indent=4)
          try:
              json.load(open(res_file))
          except Exception:
              content = []
              with open(res_file, 'r') as f:
                  for line in f:
                      content.append(line)
              content[-1] = ']'
              with open(res_file, 'w') as f:
                  for c in content:
                      f.write(c)
  
      def _coco_keypoint_results_one_category_kernel(self, data_pack):
          cat_id = data_pack['cat_id']
          keypoints = data_pack['keypoints']
          cat_results = []
  
          for img_kpts in keypoints:
              if len(img_kpts) == 0:
                  continue
  
              _key_points = np.array([img_kpts[k]['keypoints']
                                      for k in range(len(img_kpts))])
              key_points = np.zeros(
                  (_key_points.shape[0], self.kp_dim * 3), dtype=np.float)
  
              for ipt in range(self.kp_dim):
                  key_points[:, ipt * 3 + 0] = _key_points[:, ipt, 0]
                  key_points[:, ipt * 3 + 1] = _key_points[:, ipt, 1]
                  key_points[:, ipt * 3 + 2] = _key_points[:, ipt, 2]  # keypoints score.
  
              result = [{'image_id': img_kpts[k]['image'],
                         'category_id': cat_id,
                         'keypoints': list(key_points[k]),
                         'score': img_kpts[k]['score'],
                         'center': list(img_kpts[k]['center']),
                         'scale': list(img_kpts[k]['scale'])
                         } for k in range(len(img_kpts))]
              cat_results.extend(result)
  
          return cat_results
  
      def _do_python_keypoint_eval(self, res_file, res_folder):
          coco_dt = self.coco.loadRes(res_file)
          coco_eval = COCOeval(self.coco, coco_dt, 'keypoints')
          coco_eval.params.useSegm = None
          coco_eval.evaluate()
          coco_eval.accumulate()
          coco_eval.summarize()
          stats_names = ['AP', 'Ap .5', 'AP .75', 'AP (M)', 'AP (L)', 'AR', 'AR .5', 'AR .75', 'AR (M)', 'AR (L)']
  
          info_str = []
          for ind, name in enumerate(stats_names):
              info_str.append((name, coco_eval.stats[ind]))
  
          eval_file = os.path.join(res_folder, 'keypoints_val_results.pkl')
  
          with open(eval_file, 'wb') as f:
              pickle.dump(coco_eval, f, pickle.HIGHEST_PROTOCOL)
  
          print('=> coco eval results saved to %s' % eval_file)
  
          return info_str
  
  