# 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