def __init__(self, ann_file, img_prefix, data_cfg, pipeline, test_mode=False): super(TopDownCocoDataset, self).__init__(ann_file, img_prefix, data_cfg, pipeline, test_mode=test_mode) self.use_gt_bbox = data_cfg['use_gt_bbox'] self.bbox_file = data_cfg['bbox_file'] self.det_bbox_thr = data_cfg.get('det_bbox_thr', 0.0) if 'image_thr' in data_cfg: warnings.warn( 'image_thr is deprecated, ' 'please use det_bbox_thr instead', DeprecationWarning) self.det_bbox_thr = data_cfg['image_thr'] self.use_nms = data_cfg.get('use_nms', True) self.soft_nms = data_cfg['soft_nms'] self.nms_thr = data_cfg['nms_thr'] self.oks_thr = data_cfg['oks_thr'] self.vis_thr = data_cfg['vis_thr'] self.ann_info['flip_pairs'] = self._make_flip_pairs() self.ann_info['upper_body_ids'] = (0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10) self.ann_info['lower_body_ids'] = (11, 12, 13, 14, 15, 16) self.ann_info['use_different_joint_weights'] = False self.ann_info['joint_weights'] = \ np.ones((self.ann_info['num_joints'], 1), dtype=np.float32) self.body_num = 17 self.foot_num = 6 self.face_num = 68 self.left_hand_num = 21 self.right_hand_num = 21 # 'https://github.com/jin-s13/COCO-WholeBody/blob/master/' # 'evaluation/myeval_wholebody.py#L170' self.sigmas_body = [ 0.026, 0.025, 0.025, 0.035, 0.035, 0.079, 0.079, 0.072, 0.072, 0.062, 0.062, 0.107, 0.107, 0.087, 0.087, 0.089, 0.089 ] self.sigmas_foot = [0.068, 0.066, 0.066, 0.092, 0.094, 0.094] self.sigmas_face = [ 0.042, 0.043, 0.044, 0.043, 0.040, 0.035, 0.031, 0.025, 0.020, 0.023, 0.029, 0.032, 0.037, 0.038, 0.043, 0.041, 0.045, 0.013, 0.012, 0.011, 0.011, 0.012, 0.012, 0.011, 0.011, 0.013, 0.015, 0.009, 0.007, 0.007, 0.007, 0.012, 0.009, 0.008, 0.016, 0.010, 0.017, 0.011, 0.009, 0.011, 0.009, 0.007, 0.013, 0.008, 0.011, 0.012, 0.010, 0.034, 0.008, 0.008, 0.009, 0.008, 0.008, 0.007, 0.010, 0.008, 0.009, 0.009, 0.009, 0.007, 0.007, 0.008, 0.011, 0.008, 0.008, 0.008, 0.01, 0.008 ] self.sigmas_lefthand = [ 0.029, 0.022, 0.035, 0.037, 0.047, 0.026, 0.025, 0.024, 0.035, 0.018, 0.024, 0.022, 0.026, 0.017, 0.021, 0.021, 0.032, 0.02, 0.019, 0.022, 0.031 ] self.sigmas_righthand = [ 0.029, 0.022, 0.035, 0.037, 0.047, 0.026, 0.025, 0.024, 0.035, 0.018, 0.024, 0.022, 0.026, 0.017, 0.021, 0.021, 0.032, 0.02, 0.019, 0.022, 0.031 ] self.sigmas_wholebody = (self.sigmas_body + self.sigmas_foot + self.sigmas_face + self.sigmas_lefthand + self.sigmas_righthand) self.sigmas = np.array(self.sigmas_wholebody) self.coco = COCO(ann_file) 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:]) self.img_ids = self.coco.getImgIds() self.num_images = len(self.img_ids) self.id2name, self.name2id = self._get_mapping_id_name(self.coco.imgs) self.dataset_name = 'coco_wholebody' self.db = self._get_db() print(f'=> num_images: {self.num_images}') print(f'=> load {len(self.db)} samples')
class AnimalPoseDataset(AnimalBaseDataset): """Animal-Pose dataset for animal pose estimation. `Cross-domain Adaptation For Animal Pose Estimation’ ICCV'2019 More details can be found in the `paper <https://arxiv.org/abs/1908.05806>`__ . The dataset loads raw features and apply specified transforms to return a dict containing the image tensors and other information. Animal-Pose keypoint indexes:: 0: 'L_Eye', 1: 'R_Eye', 2: 'L_EarBase', 3: 'R_EarBase', 4: 'Nose', 5: 'Throat', 6: 'TailBase', 7: 'Withers', 8: 'L_F_Elbow', 9: 'R_F_Elbow', 10: 'L_B_Elbow', 11: 'R_B_Elbow', 12: 'L_F_Knee', 13: 'R_F_Knee', 14: 'L_B_Knee', 15: 'R_B_Knee', 16: 'L_F_Paw', 17: 'R_F_Paw', 18: 'L_B_Paw', 19: 'R_B_Paw' Args: ann_file (str): Path to the annotation file. img_prefix (str): Path to a directory where images are held. Default: None. data_cfg (dict): config pipeline (list[dict | callable]): A sequence of data transforms. test_mode (bool): Store True when building test or validation dataset. Default: False. """ def __init__(self, ann_file, img_prefix, data_cfg, pipeline, test_mode=False): super().__init__(ann_file, img_prefix, data_cfg, pipeline, test_mode=test_mode) self.use_gt_bbox = data_cfg['use_gt_bbox'] self.bbox_file = data_cfg['bbox_file'] self.det_bbox_thr = data_cfg.get('det_bbox_thr', 0.0) if 'image_thr' in data_cfg: warnings.warn( 'image_thr is deprecated, ' 'please use det_bbox_thr instead', DeprecationWarning) self.det_bbox_thr = data_cfg['image_thr'] self.use_nms = data_cfg.get('use_nms', True) self.soft_nms = data_cfg['soft_nms'] self.nms_thr = data_cfg['nms_thr'] self.oks_thr = data_cfg['oks_thr'] self.vis_thr = data_cfg['vis_thr'] self.ann_info['flip_pairs'] = [[0, 1], [2, 3], [8, 9], [10, 11], [12, 13], [14, 15], [16, 17], [18, 19]] self.ann_info['upper_body_ids'] = (0, 1, 2, 3, 4, 5, 7, 8, 9, 12, 13, 16, 17) self.ann_info['lower_body_ids'] = (6, 10, 11, 14, 15, 18, 19) self.ann_info['use_different_joint_weights'] = False self.ann_info['joint_weights'] = np.array( [ 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1.2, 1.2, 1.2, 1.2, 1.5, 1.5, 1.5, 1.5 ], dtype=np.float32).reshape((self.ann_info['num_joints'], 1)) # Note: The original paper did not provide enough information about # the sigmas. We modified from 'https://github.com/cocodataset/' # 'cocoapi/blob/master/PythonAPI/pycocotools/cocoeval.py#L523' self.sigmas = np.array([ .25, .25, .26, .35, .35, 1.0, 1.0, 1.0, 1.07, 1.07, 1.07, 1.07, .87, .87, .87, .87, .89, .89, .89, .89 ]) / 10.0 self.coco = COCO(ann_file) 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:]) self.img_ids = self.coco.getImgIds() self.num_images = len(self.img_ids) self.id2name, self.name2id = self._get_mapping_id_name(self.coco.imgs) self.dataset_name = 'animalpose' self.db = self._get_db() print(f'=> num_images: {self.num_images}') print(f'=> load {len(self.db)} samples') def _get_db(self): """Load dataset.""" assert self.use_gt_bbox gt_db = self._load_coco_keypoint_annotations() return gt_db def _load_coco_keypoint_annotations(self): """Ground truth bbox and keypoints.""" gt_db = [] for img_id in self.img_ids: gt_db.extend(self._load_coco_keypoint_annotation_kernel(img_id)) return gt_db def _load_coco_keypoint_annotation_kernel(self, img_id): """load annotation from COCOAPI. Note: bbox:[x1, y1, w, h] Args: img_id: coco image id Returns: dict: db entry """ img_ann = self.coco.loadImgs(img_id)[0] width = img_ann['width'] height = img_ann['height'] num_joints = self.ann_info['num_joints'] ann_ids = self.coco.getAnnIds(imgIds=img_id, iscrowd=False) objs = self.coco.loadAnns(ann_ids) # sanitize bboxes valid_objs = [] for obj in objs: if 'bbox' not in obj: continue x, y, w, h = obj['bbox'] x1 = max(0, x) y1 = max(0, y) x2 = min(width - 1, x1 + max(0, w - 1)) y2 = min(height - 1, y1 + max(0, h - 1)) if ('area' not in obj or obj['area'] > 0) and x2 > x1 and y2 > y1: obj['clean_bbox'] = [x1, y1, x2 - x1, y2 - y1] valid_objs.append(obj) objs = valid_objs bbox_id = 0 rec = [] for obj in objs: if 'keypoints' not in obj: continue if max(obj['keypoints']) == 0: continue if 'num_keypoints' in obj and obj['num_keypoints'] == 0: continue joints_3d = np.zeros((num_joints, 3), dtype=np.float32) joints_3d_visible = np.zeros((num_joints, 3), dtype=np.float32) keypoints = np.array(obj['keypoints']).reshape(-1, 3) joints_3d[:, :2] = keypoints[:, :2] joints_3d_visible[:, :2] = np.minimum(1, keypoints[:, 2:3]) center, scale = self._xywh2cs(*obj['clean_bbox'][:4]) image_file = os.path.join(self.img_prefix, self.id2name[img_id]) rec.append({ 'image_file': image_file, 'center': center, 'scale': scale, 'bbox': obj['clean_bbox'][:4], 'rotation': 0, 'joints_3d': joints_3d, 'joints_3d_visible': joints_3d_visible, 'dataset': self.dataset_name, 'bbox_score': 1, 'bbox_id': bbox_id }) bbox_id = bbox_id + 1 return rec def evaluate(self, outputs, res_folder, metric='mAP', **kwargs): """Evaluate coco keypoint results. The pose prediction results will be saved in `${res_folder}/result_keypoints.json`. Note: batch_size: N num_keypoints: K heatmap height: H heatmap width: W Args: outputs (list(dict)) :preds (np.ndarray[N,K,3]): The first two dimensions are coordinates, score is the third dimension of the array. :boxes (np.ndarray[N,6]): [center[0], center[1], scale[0] , scale[1],area, score] :image_paths (list[str]): For example, ['data/coco/val2017 /000000393226.jpg'] :heatmap (np.ndarray[N, K, H, W]): model output heatmap :bbox_id (list(int)). res_folder (str): Path of directory to save the results. metric (str | list[str]): Metric to be performed. Defaults: 'mAP'. Returns: dict: Evaluation results for evaluation metric. """ metrics = metric if isinstance(metric, list) else [metric] allowed_metrics = ['mAP'] for metric in metrics: if metric not in allowed_metrics: raise KeyError(f'metric {metric} is not supported') res_file = os.path.join(res_folder, 'result_keypoints.json') kpts = defaultdict(list) for output in outputs: preds = output['preds'] boxes = output['boxes'] image_paths = output['image_paths'] bbox_ids = output['bbox_ids'] batch_size = len(image_paths) for i in range(batch_size): image_id = self.name2id[image_paths[i][len(self.img_prefix):]] kpts[image_id].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) # rescoring and oks nms num_joints = self.ann_info['num_joints'] vis_thr = self.vis_thr oks_thr = self.oks_thr valid_kpts = [] for image_id in kpts.keys(): img_kpts = kpts[image_id] for n_p in img_kpts: box_score = n_p['score'] kpt_score = 0 valid_num = 0 for n_jt in range(0, num_joints): t_s = n_p['keypoints'][n_jt][2] if t_s > vis_thr: 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 if self.use_nms: nms = soft_oks_nms if self.soft_nms else oks_nms keep = nms(list(img_kpts), oks_thr, sigmas=self.sigmas) valid_kpts.append([img_kpts[_keep] for _keep in keep]) else: valid_kpts.append(img_kpts) self._write_coco_keypoint_results(valid_kpts, res_file) info_str = self._do_python_keypoint_eval(res_file) name_value = OrderedDict(info_str) return name_value def _write_coco_keypoint_results(self, keypoints, res_file): """Write results into a json 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) def _coco_keypoint_results_one_category_kernel(self, data_pack): """Get coco keypoint results.""" 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_kpt['keypoints'] for img_kpt in img_kpts]) key_points = _key_points.reshape(-1, self.ann_info['num_joints'] * 3) result = [{ 'image_id': img_kpt['image_id'], 'category_id': cat_id, 'keypoints': key_point.tolist(), 'score': float(img_kpt['score']), 'center': img_kpt['center'].tolist(), 'scale': img_kpt['scale'].tolist() } for img_kpt, key_point in zip(img_kpts, key_points)] cat_results.extend(result) return cat_results def _do_python_keypoint_eval(self, res_file): """Keypoint evaluation using COCOAPI.""" coco_det = self.coco.loadRes(res_file) coco_eval = COCOeval(self.coco, coco_det, 'keypoints', self.sigmas) 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 = list(zip(stats_names, coco_eval.stats)) return info_str def _sort_and_unique_bboxes(self, kpts, key='bbox_id'): """sort kpts and remove the repeated ones.""" for img_id, persons in kpts.items(): num = len(persons) kpts[img_id] = sorted(kpts[img_id], key=lambda x: x[key]) for i in range(num - 1, 0, -1): if kpts[img_id][i][key] == kpts[img_id][i - 1][key]: del kpts[img_id][i] return kpts
class TopDownCocoWholeBodyDataset(TopDownCocoDataset): """CocoWholeBodyDataset dataset for top-down pose estimation. `Whole-Body Human Pose Estimation in the Wild' ECCV'2020 More details can be found in the `paper <https://arxiv.org/abs/2007.11858>`__ . The dataset loads raw features and apply specified transforms to return a dict containing the image tensors and other information. In total, we have 133 keypoints for wholebody pose estimation. COCO-WholeBody keypoint indexes:: 0-16: 17 body keypoints 17-22: 6 foot keypoints 23-90: 68 face keypoints 91-132: 42 hand keypoints Args: ann_file (str): Path to the annotation file. img_prefix (str): Path to a directory where images are held. Default: None. data_cfg (dict): config pipeline (list[dict | callable]): A sequence of data transforms. test_mode (bool): Store True when building test or validation dataset. Default: False. """ def __init__(self, ann_file, img_prefix, data_cfg, pipeline, test_mode=False): super(TopDownCocoDataset, self).__init__(ann_file, img_prefix, data_cfg, pipeline, test_mode=test_mode) self.use_gt_bbox = data_cfg['use_gt_bbox'] self.bbox_file = data_cfg['bbox_file'] self.det_bbox_thr = data_cfg.get('det_bbox_thr', 0.0) if 'image_thr' in data_cfg: warnings.warn( 'image_thr is deprecated, ' 'please use det_bbox_thr instead', DeprecationWarning) self.det_bbox_thr = data_cfg['image_thr'] self.use_nms = data_cfg.get('use_nms', True) self.soft_nms = data_cfg['soft_nms'] self.nms_thr = data_cfg['nms_thr'] self.oks_thr = data_cfg['oks_thr'] self.vis_thr = data_cfg['vis_thr'] self.ann_info['flip_pairs'] = self._make_flip_pairs() self.ann_info['upper_body_ids'] = (0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10) self.ann_info['lower_body_ids'] = (11, 12, 13, 14, 15, 16) self.ann_info['use_different_joint_weights'] = False self.ann_info['joint_weights'] = \ np.ones((self.ann_info['num_joints'], 1), dtype=np.float32) self.body_num = 17 self.foot_num = 6 self.face_num = 68 self.left_hand_num = 21 self.right_hand_num = 21 # 'https://github.com/jin-s13/COCO-WholeBody/blob/master/' # 'evaluation/myeval_wholebody.py#L170' self.sigmas_body = [ 0.026, 0.025, 0.025, 0.035, 0.035, 0.079, 0.079, 0.072, 0.072, 0.062, 0.062, 0.107, 0.107, 0.087, 0.087, 0.089, 0.089 ] self.sigmas_foot = [0.068, 0.066, 0.066, 0.092, 0.094, 0.094] self.sigmas_face = [ 0.042, 0.043, 0.044, 0.043, 0.040, 0.035, 0.031, 0.025, 0.020, 0.023, 0.029, 0.032, 0.037, 0.038, 0.043, 0.041, 0.045, 0.013, 0.012, 0.011, 0.011, 0.012, 0.012, 0.011, 0.011, 0.013, 0.015, 0.009, 0.007, 0.007, 0.007, 0.012, 0.009, 0.008, 0.016, 0.010, 0.017, 0.011, 0.009, 0.011, 0.009, 0.007, 0.013, 0.008, 0.011, 0.012, 0.010, 0.034, 0.008, 0.008, 0.009, 0.008, 0.008, 0.007, 0.010, 0.008, 0.009, 0.009, 0.009, 0.007, 0.007, 0.008, 0.011, 0.008, 0.008, 0.008, 0.01, 0.008 ] self.sigmas_lefthand = [ 0.029, 0.022, 0.035, 0.037, 0.047, 0.026, 0.025, 0.024, 0.035, 0.018, 0.024, 0.022, 0.026, 0.017, 0.021, 0.021, 0.032, 0.02, 0.019, 0.022, 0.031 ] self.sigmas_righthand = [ 0.029, 0.022, 0.035, 0.037, 0.047, 0.026, 0.025, 0.024, 0.035, 0.018, 0.024, 0.022, 0.026, 0.017, 0.021, 0.021, 0.032, 0.02, 0.019, 0.022, 0.031 ] self.sigmas_wholebody = (self.sigmas_body + self.sigmas_foot + self.sigmas_face + self.sigmas_lefthand + self.sigmas_righthand) self.sigmas = np.array(self.sigmas_wholebody) self.coco = COCO(ann_file) 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:]) self.img_ids = self.coco.getImgIds() self.num_images = len(self.img_ids) self.id2name, self.name2id = self._get_mapping_id_name(self.coco.imgs) self.dataset_name = 'coco_wholebody' self.db = self._get_db() print(f'=> num_images: {self.num_images}') print(f'=> load {len(self.db)} samples') @staticmethod def _make_flip_pairs(): body = [[1, 2], [3, 4], [5, 6], [7, 8], [9, 10], [11, 12], [13, 14], [15, 16]] foot = [[17, 20], [18, 21], [19, 22]] face = [[23, 39], [24, 38], [25, 37], [26, 36], [27, 35], [28, 34], [29, 33], [30, 32], [40, 49], [41, 48], [42, 47], [43, 46], [44, 45], [54, 58], [55, 57], [59, 68], [60, 67], [61, 66], [62, 65], [63, 70], [64, 69], [71, 77], [72, 76], [73, 75], [78, 82], [79, 81], [83, 87], [84, 86], [88, 90]] hand = [[91, 112], [92, 113], [93, 114], [94, 115], [95, 116], [96, 117], [97, 118], [98, 119], [99, 120], [100, 121], [101, 122], [102, 123], [103, 124], [104, 125], [105, 126], [106, 127], [107, 128], [108, 129], [109, 130], [110, 131], [111, 132]] return body + foot + face + hand def _load_coco_keypoint_annotation_kernel(self, img_id): """load annotation from COCOAPI. Note: bbox:[x1, y1, w, h] Args: img_id: coco image id Returns: dict: db entry """ img_ann = self.coco.loadImgs(img_id)[0] width = img_ann['width'] height = img_ann['height'] num_joints = self.ann_info['num_joints'] ann_ids = self.coco.getAnnIds(imgIds=img_id, iscrowd=False) objs = self.coco.loadAnns(ann_ids) # sanitize bboxes valid_objs = [] for obj in objs: x, y, w, h = obj['bbox'] x1 = max(0, x) y1 = max(0, y) x2 = min(width - 1, x1 + max(0, w - 1)) y2 = min(height - 1, y1 + max(0, h - 1)) if ('area' not in obj or obj['area'] > 0) and x2 > x1 and y2 > y1: obj['clean_bbox'] = [x1, y1, x2 - x1, y2 - y1] valid_objs.append(obj) objs = valid_objs rec = [] bbox_id = 0 for obj in objs: if max(obj['keypoints']) == 0: continue joints_3d = np.zeros((num_joints, 3), dtype=np.float32) joints_3d_visible = np.zeros((num_joints, 3), dtype=np.float32) keypoints = np.array(obj['keypoints'] + obj['foot_kpts'] + obj['face_kpts'] + obj['lefthand_kpts'] + obj['righthand_kpts']).reshape(-1, 3) joints_3d[:, :2] = keypoints[:, :2] joints_3d_visible[:, :2] = np.minimum(1, keypoints[:, 2:3] > 0) center, scale = self._xywh2cs(*obj['clean_bbox'][:4]) image_file = os.path.join(self.img_prefix, self.id2name[img_id]) rec.append({ 'image_file': image_file, 'center': center, 'scale': scale, 'rotation': 0, 'joints_3d': joints_3d, 'joints_3d_visible': joints_3d_visible, 'dataset': self.dataset_name, 'bbox_score': 1, 'bbox_id': bbox_id }) bbox_id = bbox_id + 1 return rec def _coco_keypoint_results_one_category_kernel(self, data_pack): """Get coco keypoint results.""" 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_kpt['keypoints'] for img_kpt in img_kpts]) key_points = _key_points.reshape(-1, self.ann_info['num_joints'] * 3) cuts = np.cumsum([ 0, self.body_num, self.foot_num, self.face_num, self.left_hand_num, self.right_hand_num ]) * 3 result = [{ 'image_id': img_kpt['image_id'], 'category_id': cat_id, 'keypoints': key_point[cuts[0]:cuts[1]].tolist(), 'foot_kpts': key_point[cuts[1]:cuts[2]].tolist(), 'face_kpts': key_point[cuts[2]:cuts[3]].tolist(), 'lefthand_kpts': key_point[cuts[3]:cuts[4]].tolist(), 'righthand_kpts': key_point[cuts[4]:cuts[5]].tolist(), 'score': float(img_kpt['score']), 'center': img_kpt['center'].tolist(), 'scale': img_kpt['scale'].tolist() } for img_kpt, key_point in zip(img_kpts, key_points)] cat_results.extend(result) return cat_results def _do_python_keypoint_eval(self, res_file): """Keypoint evaluation using COCOAPI.""" coco_det = self.coco.loadRes(res_file) coco_eval = COCOeval(self.coco, coco_det, 'keypoints_body', np.array(self.sigmas_body), 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_foot', np.array(self.sigmas_foot), 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_face', np.array(self.sigmas_face), 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_lefthand', np.array(self.sigmas_lefthand), 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_righthand', np.array(self.sigmas_righthand), 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', np.array(self.sigmas_wholebody), use_area=True) 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 = list(zip(stats_names, coco_eval.stats)) return info_str
def cocoapi_eval(jsonfile, style, coco_gt=None, anno_file=None, max_dets=(100, 300, 1000), classwise=False, sigmas=None, use_area=True): """ Args: jsonfile (str): Evaluation json file, eg: bbox.json, mask.json. style (str): COCOeval style, can be `bbox` , `segm` , `proposal`, `keypoints` and `keypoints_crowd`. coco_gt (str): Whether to load COCOAPI through anno_file, eg: coco_gt = COCO(anno_file) anno_file (str): COCO annotations file. max_dets (tuple): COCO evaluation maxDets. classwise (bool): Whether per-category AP and draw P-R Curve or not. sigmas (nparray): keypoint labelling sigmas. use_area (bool): If gt annotations (eg. CrowdPose, AIC) do not have 'area', please set use_area=False. """ assert coco_gt != None or anno_file != None if style == 'keypoints_crowd': #please install xtcocotools==1.6 from xtcocotools.coco import COCO from xtcocotools.cocoeval import COCOeval else: from pycocotools.coco import COCO from pycocotools.cocoeval import COCOeval if coco_gt == None: coco_gt = COCO(anno_file) logger.info("Start evaluate...") coco_dt = coco_gt.loadRes(jsonfile) if style == 'proposal': coco_eval = COCOeval(coco_gt, coco_dt, 'bbox') coco_eval.params.useCats = 0 coco_eval.params.maxDets = list(max_dets) elif style == 'keypoints_crowd': coco_eval = COCOeval(coco_gt, coco_dt, style, sigmas, use_area) else: coco_eval = COCOeval(coco_gt, coco_dt, style) coco_eval.evaluate() coco_eval.accumulate() coco_eval.summarize() if classwise: # Compute per-category AP and PR curve try: from terminaltables import AsciiTable except Exception as e: logger.error( 'terminaltables not found, plaese install terminaltables. ' 'for example: `pip install terminaltables`.') raise e precisions = coco_eval.eval['precision'] cat_ids = coco_gt.getCatIds() # precision: (iou, recall, cls, area range, max dets) assert len(cat_ids) == precisions.shape[2] results_per_category = [] for idx, catId in enumerate(cat_ids): # area range index 0: all area ranges # max dets index -1: typically 100 per image nm = coco_gt.loadCats(catId)[0] precision = precisions[:, :, idx, 0, -1] precision = precision[precision > -1] if precision.size: ap = np.mean(precision) else: ap = float('nan') results_per_category.append( (str(nm["name"]), '{:0.3f}'.format(float(ap)))) pr_array = precisions[0, :, idx, 0, 2] recall_array = np.arange(0.0, 1.01, 0.01) draw_pr_curve(pr_array, recall_array, out_dir=style + '_pr_curve', file_name='{}_precision_recall_curve.jpg'.format( nm["name"])) num_columns = min(6, len(results_per_category) * 2) results_flatten = list(itertools.chain(*results_per_category)) headers = ['category', 'AP'] * (num_columns // 2) results_2d = itertools.zip_longest( *[results_flatten[i::num_columns] for i in range(num_columns)]) table_data = [headers] table_data += [result for result in results_2d] table = AsciiTable(table_data) logger.info('Per-category of {} AP: \n{}'.format(style, table.table)) logger.info("per-category PR curve has output to {} folder.".format( style + '_pr_curve')) # flush coco evaluation result sys.stdout.flush() return coco_eval.stats
def test_bottomup_pipeline(): data_prefix = 'tests/data/coco/' ann_file = osp.join(data_prefix, 'test_coco.json') coco = COCO(ann_file) ann_info = {} ann_info['flip_pairs'] = [[1, 2], [3, 4], [5, 6], [7, 8], [9, 10], [11, 12], [13, 14], [15, 16]] ann_info['flip_index'] = [ 0, 2, 1, 4, 3, 6, 5, 8, 7, 10, 9, 12, 11, 14, 13, 16, 15 ] ann_info['use_different_joint_weights'] = False ann_info['joint_weights'] = np.array([ 1., 1., 1., 1., 1., 1., 1., 1.2, 1.2, 1.5, 1.5, 1., 1., 1.2, 1.2, 1.5, 1.5 ], dtype=np.float32).reshape((17, 1)) ann_info['image_size'] = np.array(512) ann_info['heatmap_size'] = np.array([128, 256]) ann_info['num_joints'] = 17 ann_info['num_scales'] = 2 ann_info['scale_aware_sigma'] = False ann_ids = coco.getAnnIds(785) anno = coco.loadAnns(ann_ids) mask = _get_mask(coco, anno, 785) anno = [ obj for obj in anno if obj['iscrowd'] == 0 or obj['num_keypoints'] > 0 ] joints = _get_joints(anno, ann_info, False) mask_list = [mask.copy() for _ in range(ann_info['num_scales'])] joints_list = [joints.copy() for _ in range(ann_info['num_scales'])] results = {} results['dataset'] = 'coco' results['image_file'] = osp.join(data_prefix, '000000000785.jpg') results['mask'] = mask_list results['joints'] = joints_list results['ann_info'] = ann_info transform = LoadImageFromFile() results = transform(copy.deepcopy(results)) assert results['img'].shape == (425, 640, 3) # test HorizontalFlip random_horizontal_flip = BottomUpRandomFlip(flip_prob=1.) results_horizontal_flip = random_horizontal_flip(copy.deepcopy(results)) assert _check_flip(results['img'], results_horizontal_flip['img']) random_horizontal_flip = BottomUpRandomFlip(flip_prob=0.) results_horizontal_flip = random_horizontal_flip(copy.deepcopy(results)) assert (results['img'] == results_horizontal_flip['img']).all() results_copy = copy.deepcopy(results) results_copy['mask'] = mask_list[0] with pytest.raises(AssertionError): results_horizontal_flip = random_horizontal_flip( copy.deepcopy(results_copy)) results_copy = copy.deepcopy(results) results_copy['joints'] = joints_list[0] with pytest.raises(AssertionError): results_horizontal_flip = random_horizontal_flip( copy.deepcopy(results_copy)) results_copy = copy.deepcopy(results) results_copy['joints'] = joints_list[:1] with pytest.raises(AssertionError): results_horizontal_flip = random_horizontal_flip( copy.deepcopy(results_copy)) results_copy = copy.deepcopy(results) results_copy['mask'] = mask_list[:1] with pytest.raises(AssertionError): results_horizontal_flip = random_horizontal_flip( copy.deepcopy(results_copy)) # test TopDownAffine random_affine_transform = BottomUpRandomAffine(30, [0.75, 1.5], 'short', 0) results_affine_transform = random_affine_transform(copy.deepcopy(results)) assert results_affine_transform['img'].shape == (512, 512, 3) random_affine_transform = BottomUpRandomAffine(30, [0.75, 1.5], 'short', 40) results_affine_transform = random_affine_transform(copy.deepcopy(results)) assert results_affine_transform['img'].shape == (512, 512, 3) results_copy = copy.deepcopy(results) results_copy['ann_info']['scale_aware_sigma'] = True joints = _get_joints(anno, results_copy['ann_info'], False) results_copy['joints'] = \ [joints.copy() for _ in range(results_copy['ann_info']['num_scales'])] results_affine_transform = random_affine_transform(results_copy) assert results_affine_transform['img'].shape == (512, 512, 3) results_copy = copy.deepcopy(results) results_copy['mask'] = mask_list[0] with pytest.raises(AssertionError): results_horizontal_flip = random_affine_transform( copy.deepcopy(results_copy)) results_copy = copy.deepcopy(results) results_copy['joints'] = joints_list[0] with pytest.raises(AssertionError): results_horizontal_flip = random_affine_transform( copy.deepcopy(results_copy)) results_copy = copy.deepcopy(results) results_copy['joints'] = joints_list[:1] with pytest.raises(AssertionError): results_horizontal_flip = random_affine_transform( copy.deepcopy(results_copy)) results_copy = copy.deepcopy(results) results_copy['mask'] = mask_list[:1] with pytest.raises(AssertionError): results_horizontal_flip = random_affine_transform( copy.deepcopy(results_copy)) random_affine_transform = BottomUpRandomAffine(30, [0.75, 1.5], 'long', 40) results_affine_transform = random_affine_transform(copy.deepcopy(results)) assert results_affine_transform['img'].shape == (512, 512, 3) with pytest.raises(ValueError): random_affine_transform = BottomUpRandomAffine(30, [0.75, 1.5], 'short-long', 40) results_affine_transform = random_affine_transform( copy.deepcopy(results)) # test BottomUpGenerateTarget generate_multi_target = BottomUpGenerateTarget(2, 30) results_generate_multi_target = generate_multi_target( copy.deepcopy(results)) assert 'targets' in results_generate_multi_target assert len(results_generate_multi_target['targets'] ) == results['ann_info']['num_scales'] # test BottomUpGetImgSize when W > H get_multi_scale_size = BottomUpGetImgSize([1]) results_get_multi_scale_size = get_multi_scale_size(copy.deepcopy(results)) assert 'test_scale_factor' in results_get_multi_scale_size['ann_info'] assert 'base_size' in results_get_multi_scale_size['ann_info'] assert 'center' in results_get_multi_scale_size['ann_info'] assert 'scale' in results_get_multi_scale_size['ann_info'] assert results_get_multi_scale_size['ann_info']['base_size'][1] == 512 # test BottomUpResizeAlign transforms = [ dict(type='ToTensor'), dict( type='NormalizeTensor', mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), ] resize_align_multi_scale = BottomUpResizeAlign(transforms=transforms) results_copy = copy.deepcopy(results_get_multi_scale_size) results_resize_align_multi_scale = resize_align_multi_scale(results_copy) assert 'aug_data' in results_resize_align_multi_scale['ann_info'] # test BottomUpGetImgSize when W < H results_copy = copy.deepcopy(results) results_copy['img'] = np.random.rand(640, 425, 3) results_get_multi_scale_size = get_multi_scale_size(results_copy) assert results_get_multi_scale_size['ann_info']['base_size'][0] == 512
class HandBaseDataset(Dataset, metaclass=ABCMeta): """Base class for hand datasets. All hand datasets should subclass it. All subclasses should overwrite: Methods:`_get_db`, 'evaluate' Args: ann_file (str): Path to the annotation file. img_prefix (str): Path to a directory where images are held. Default: None. data_cfg (dict): config pipeline (list[dict | callable]): A sequence of data transforms. test_mode (bool): Store True when building test or validation dataset. Default: False. """ def __init__(self, ann_file, img_prefix, data_cfg, pipeline, test_mode=False): self.image_info = {} self.ann_info = {} self.annotations_path = ann_file self.img_prefix = img_prefix self.pipeline = pipeline self.test_mode = test_mode self.ann_info['image_size'] = np.array(data_cfg['image_size']) self.ann_info['heatmap_size'] = np.array(data_cfg['heatmap_size']) self.ann_info['num_joints'] = data_cfg['num_joints'] self.ann_info['flip_pairs'] = [] self.ann_info['inference_channel'] = data_cfg['inference_channel'] self.ann_info['num_output_channels'] = data_cfg['num_output_channels'] self.ann_info['dataset_channel'] = data_cfg['dataset_channel'] self.coco = COCO(ann_file) self.img_ids = self.coco.getImgIds() self.num_images = len(self.img_ids) self.id2name, self.name2id = self._get_mapping_id_name(self.coco.imgs) self.db = [] self.pipeline = Compose(self.pipeline) def _get_mapping_id_name(self, imgs): """ Args: imgs (dict): dict of image info. Returns: tuple: Image name & id mapping dicts. - id2name (dict): Mapping image id to name. - name2id (dict): Mapping image name to id. """ id2name = {} name2id = {} for image_id, image in imgs.items(): file_name = image['file_name'] id2name[image_id] = file_name name2id[file_name] = image_id return id2name, name2id def _xywh2cs(self, x, y, w, h, padding=1.25): """This encodes bbox(x,y,w,w) into (center, scale) Args: x, y, w, h Returns: center (np.ndarray[float32](2,)): center of the bbox (x, y). scale (np.ndarray[float32](2,)): scale of the bbox w & h. """ aspect_ratio = self.ann_info['image_size'][0] / self.ann_info[ 'image_size'][1] center = np.array([x + w * 0.5, y + h * 0.5], dtype=np.float32) if (not self.test_mode) and np.random.rand() < 0.3: center += 0.4 * (np.random.rand(2) - 0.5) * [w, h] if w > aspect_ratio * h: h = w * 1.0 / aspect_ratio elif w < aspect_ratio * h: w = h * aspect_ratio # pixel std is 200.0 scale = np.array([w / 200.0, h / 200.0], dtype=np.float32) scale = scale * padding return center, scale @abstractmethod def _get_db(self): """Load dataset.""" raise NotImplementedError @abstractmethod def evaluate(self, cfg, preds, output_dir, *args, **kwargs): """Evaluate keypoint results.""" raise NotImplementedError def _write_keypoint_results(self, keypoints, res_file): """Write results into a json file.""" with open(res_file, 'w') as f: json.dump(keypoints, f, sort_keys=True, indent=4) def _report_metric(self, res_file, metrics, pck_thr=0.2, pckh_thr=0.7, auc_nor=30): """Keypoint evaluation. Args: res_file (str): Json file stored prediction results. metrics (str | list[str]): Metric to be performed. Options: 'PCK', 'PCKh', 'AUC', 'EPE'. pck_thr (float): PCK threshold, default as 0.2. pckh_thr (float): PCKh threshold, default as 0.7. auc_nor (float): AUC normalization factor, default as 30 pixel. Returns: dict: Evaluation results for evaluation metric. """ info_str = [] with open(res_file, 'r') as fin: preds = json.load(fin) assert len(preds) == len(self.db) outputs = [] gts = [] masks = [] threshold_bbox = [] threshold_head_box = [] for pred, item in zip(preds, self.db): outputs.append(np.array(pred['keypoints'])[:, :-1]) gts.append(np.array(item['joints_3d'])[:, :-1]) masks.append((np.array(item['joints_3d_visible'])[:, 0]) > 0) if 'PCK' in metrics: bbox = np.array(item['bbox']) bbox_thr = np.max(bbox[2:]) threshold_bbox.append(np.array([bbox_thr, bbox_thr])) if 'PCKh' in metrics: head_box_thr = item['head_size'] threshold_head_box.append( np.array([head_box_thr, head_box_thr])) outputs = np.array(outputs) gts = np.array(gts) masks = np.array(masks) threshold_bbox = np.array(threshold_bbox) threshold_head_box = np.array(threshold_head_box) if 'PCK' in metrics: _, pck, _ = keypoint_pck_accuracy(outputs, gts, masks, pck_thr, threshold_bbox) info_str.append(('PCK', pck)) if 'PCKh' in metrics: _, pckh, _ = keypoint_pck_accuracy(outputs, gts, masks, pckh_thr, threshold_head_box) info_str.append(('PCKh', pckh)) if 'AUC' in metrics: info_str.append(('AUC', keypoint_auc(outputs, gts, masks, auc_nor))) if 'EPE' in metrics: info_str.append(('EPE', keypoint_epe(outputs, gts, masks))) return info_str def __len__(self): """Get the size of the dataset.""" return len(self.db) def __getitem__(self, idx): """Get the sample given index.""" results = copy.deepcopy(self.db[idx]) results['ann_info'] = self.ann_info return self.pipeline(results)
class TopDownEnvisatCocoDataset(TopDownBaseDataset): """CocoDataset dataset for top-down pose estimation. `Microsoft COCO: Common Objects in Context' ECCV'2014 More details can be found in the `paper <https://arxiv.org/abs/1405.0312>`__ . The dataset loads raw features and apply specified transforms to return a dict containing the image tensors and other information. COCO keypoint indexes:: 0: 'nose', 1: 'left_eye', 2: 'right_eye', 3: 'left_ear', 4: 'right_ear', 5: 'left_shoulder', 6: 'right_shoulder', 7: 'left_elbow', 8: 'right_elbow', 9: 'left_wrist', 10: 'right_wrist', 11: 'left_hip', 12: 'right_hip', 13: 'left_knee', 14: 'right_knee', 15: 'left_ankle', 16: 'right_ankle' Args: ann_file (str): Path to the annotation file. img_prefix (str): Path to a directory where images are held. Default: None. data_cfg (dict): config pipeline (list[dict | callable]): A sequence of data transforms. test_mode (bool): Store True when building test or validation dataset. Default: False. """ def __init__(self, ann_file, img_prefix, data_cfg, pipeline, test_mode=False): super().__init__(ann_file, img_prefix, data_cfg, pipeline, test_mode=test_mode) self.use_gt_bbox = data_cfg['use_gt_bbox'] self.bbox_file = data_cfg['bbox_file'] self.det_bbox_thr = data_cfg.get('det_bbox_thr', 0.0) if 'image_thr' in data_cfg: warnings.warn( 'image_thr is deprecated, ' 'please use det_bbox_thr instead', DeprecationWarning) self.det_bbox_thr = data_cfg['image_thr'] self.use_nms = data_cfg.get('use_nms', True) self.soft_nms = data_cfg['soft_nms'] self.nms_thr = data_cfg['nms_thr'] self.oks_thr = data_cfg['oks_thr'] self.vis_thr = data_cfg['vis_thr'] self.ann_info['flip_pairs'] = [] self.ann_info['upper_body_ids'] = () self.ann_info['lower_body_ids'] = () self.ann_info['use_different_joint_weights'] = False self.ann_info['joint_weights'] = np.ones( (self.ann_info['num_joints'], 1), dtype=np.float32) # 'https://github.com/cocodataset/cocoapi/blob/master/PythonAPI/' # 'pycocotools/cocoeval.py#L523' self.sigmas = np.ones((12, )) / 10.0 self.coco = COCO(ann_file) 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:]) self.img_ids = self.coco.getImgIds() self.num_images = len(self.img_ids) self.id2name, self.name2id = self._get_mapping_id_name(self.coco.imgs) self.dataset_name = 'coco' self.db = self._get_db() print(f'=> num_images: {self.num_images}') print(f'=> load {len(self.db)} samples') @staticmethod def _get_mapping_id_name(imgs): """ Args: imgs (dict): dict of image info. Returns: tuple: Image name & id mapping dicts. - id2name (dict): Mapping image id to name. - name2id (dict): Mapping image name to id. """ id2name = {} name2id = {} for image_id, image in imgs.items(): file_name = image['file_name'] id2name[image_id] = file_name name2id[file_name] = image_id return id2name, name2id def _get_db(self): """Load dataset.""" if (not self.test_mode) or self.use_gt_bbox: # use ground truth bbox gt_db = self._load_coco_keypoint_annotations() else: # use bbox from detection gt_db = self._load_coco_person_detection_results() return gt_db def _load_coco_keypoint_annotations(self): """Ground truth bbox and keypoints.""" gt_db = [] for img_id in self.img_ids: gt_db.extend(self._load_coco_keypoint_annotation_kernel(img_id)) return gt_db def _load_coco_keypoint_annotation_kernel(self, img_id): """load annotation from COCOAPI. Note: bbox:[x1, y1, w, h] Args: img_id: coco image id Returns: dict: db entry """ img_ann = self.coco.loadImgs(img_id)[0] width = img_ann['width'] height = img_ann['height'] num_joints = self.ann_info['num_joints'] ann_ids = self.coco.getAnnIds(imgIds=img_id, iscrowd=False) objs = self.coco.loadAnns(ann_ids) # sanitize bboxes valid_objs = [] for obj in objs: if 'bbox' not in obj: continue x, y, w, h = obj['bbox'] x1 = max(0, x) y1 = max(0, y) x2 = min(width - 1, x1 + max(0, w - 1)) y2 = min(height - 1, y1 + max(0, h - 1)) if ('area' not in obj or obj['area'] > 0) and x2 > x1 and y2 > y1: obj['clean_bbox'] = [x1, y1, x2 - x1, y2 - y1] valid_objs.append(obj) objs = valid_objs bbox_id = 0 rec = [] for obj in objs: if 'keypoints' not in obj: continue if max(obj['keypoints']) == 0: continue if 'num_keypoints' in obj and obj['num_keypoints'] == 0: continue joints_3d = np.zeros((num_joints, 3), dtype=np.float32) joints_3d_visible = np.zeros((num_joints, 3), dtype=np.float32) keypoints = np.array(obj['keypoints']).reshape(-1, 3) joints_3d[:, :2] = keypoints[:, :2] joints_3d_visible[:, :2] = np.minimum(1, keypoints[:, 2:3]) center, scale = self._xywh2cs(*obj['clean_bbox'][:4]) image_file = os.path.join(self.img_prefix, self._get_subfolder_name(img_id), self.id2name[img_id]) rec.append({ 'image_file': image_file, 'center': center, 'scale': scale, 'bbox': obj['clean_bbox'][:4], 'rotation': 0, 'joints_3d': joints_3d, 'joints_3d_visible': joints_3d_visible, 'dataset': self.dataset_name, 'bbox_score': 1, 'bbox_id': bbox_id }) bbox_id = bbox_id + 1 return rec def _get_subfolder_name(self, index): #TODO IMPROVE THIS """ example: images / train2017 / 000000119993.jpg """ #For Envisat with 1000 images per folder n_subfolder = 1000 idx_range = (int(np.floor(index / n_subfolder)) * n_subfolder, int(np.ceil(index / n_subfolder)) * n_subfolder) if idx_range[0] == idx_range[1]: idx_range = (idx_range[0], (idx_range[1] + n_subfolder - 1)) else: idx_range = (idx_range[0], (idx_range[1] - 1)) subfolder = f"{idx_range[0]}-{idx_range[1]}" # file_name = self.img_prefix + format(index)+'.jpg' # prefix = self.image_set # image_path = os.path.join( # self.root, prefix,subfolder, file_name) return subfolder def _xywh2cs(self, x, y, w, h): """This encodes bbox(x,y,w,w) into (center, scale) Args: x, y, w, h Returns: tuple: A tuple containing center and scale. - center (np.ndarray[float32](2,)): center of the bbox (x, y). - scale (np.ndarray[float32](2,)): scale of the bbox w & h. """ aspect_ratio = self.ann_info['image_size'][0] / self.ann_info[ 'image_size'][1] center = np.array([x + w * 0.5, y + h * 0.5], dtype=np.float32) if (not self.test_mode) and np.random.rand() < 0.3: center += 0.4 * (np.random.rand(2) - 0.5) * [w, h] if w > aspect_ratio * h: h = w * 1.0 / aspect_ratio elif w < aspect_ratio * h: w = h * aspect_ratio # pixel std is 200.0 scale = np.array([w / 200.0, h / 200.0], dtype=np.float32) # padding to include proper amount of context scale = scale * 1.25 return center, scale def _load_coco_person_detection_results(self): """Load coco person detection results.""" num_joints = self.ann_info['num_joints'] all_boxes = None with open(self.bbox_file, 'r') as f: all_boxes = json.load(f) if not all_boxes: raise ValueError('=> Load %s fail!' % self.bbox_file) print(f'=> Total boxes: {len(all_boxes)}') kpt_db = [] bbox_id = 0 for det_res in all_boxes: if det_res['category_id'] != 1: continue img_id = det_res['image_id'] image_file = os.path.join(self.img_prefix, self._get_subfolder_name(img_id), self.id2name[img_id]) box = det_res['bbox'] score = det_res['score'] if score < self.det_bbox_thr: continue center, scale = self._xywh2cs(*box[:4]) joints_3d = np.zeros((num_joints, 3), dtype=np.float32) joints_3d_visible = np.ones((num_joints, 3), dtype=np.float32) kpt_db.append({ 'image_file': image_file, 'center': center, 'scale': scale, 'rotation': 0, 'bbox': box[:4], 'bbox_score': score, 'dataset': self.dataset_name, 'joints_3d': joints_3d, 'joints_3d_visible': joints_3d_visible, 'bbox_id': bbox_id }) bbox_id = bbox_id + 1 print(f'=> Total boxes after filter ' f'low score@{self.det_bbox_thr}: {bbox_id}') return kpt_db def evaluate(self, outputs, res_folder, metric='mAP', **kwargs): """Evaluate coco keypoint results. The pose prediction results will be saved in `${res_folder}/result_keypoints.json`. Note: batch_size: N num_keypoints: K heatmap height: H heatmap width: W Args: outputs (list(dict)) :preds (np.ndarray[N,K,3]): The first two dimensions are coordinates, score is the third dimension of the array. :boxes (np.ndarray[N,6]): [center[0], center[1], scale[0] , scale[1],area, score] :image_paths (list[str]): For example, ['data/coco/val2017 /000000393226.jpg'] :heatmap (np.ndarray[N, K, H, W]): model output heatmap :bbox_id (list(int)). res_folder (str): Path of directory to save the results. metric (str | list[str]): Metric to be performed. Defaults: 'mAP'. Returns: dict: Evaluation results for evaluation metric. """ metrics = metric if isinstance(metric, list) else [metric] allowed_metrics = ['mAP'] for metric in metrics: if metric not in allowed_metrics: raise KeyError(f'metric {metric} is not supported') res_file = os.path.join(res_folder, 'result_keypoints.json') kpts = defaultdict(list) for output in outputs: preds = output['preds'] boxes = output['boxes'] image_paths = output['image_paths'] bbox_ids = output['bbox_ids'] batch_size = len(image_paths) for i in range(batch_size): image_id = self.name2id[(image_paths[i][0:]).split("/")[-1]] kpts[image_id].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) # rescoring and oks nms num_joints = self.ann_info['num_joints'] vis_thr = self.vis_thr oks_thr = self.oks_thr valid_kpts = [] for image_id in kpts.keys(): img_kpts = kpts[image_id] for n_p in img_kpts: box_score = n_p['score'] kpt_score = 0 valid_num = 0 for n_jt in range(0, num_joints): t_s = n_p['keypoints'][n_jt][2] if t_s > vis_thr: 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 if self.use_nms: nms = soft_oks_nms if self.soft_nms else oks_nms keep = nms(list(img_kpts), oks_thr, sigmas=self.sigmas) valid_kpts.append([img_kpts[_keep] for _keep in keep]) else: valid_kpts.append(img_kpts) self._write_coco_keypoint_results(valid_kpts, res_file) info_str = self._do_python_keypoint_eval(res_file) name_value = OrderedDict(info_str) return name_value def _write_coco_keypoint_results(self, keypoints, res_file): """Write results into a json 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) def _coco_keypoint_results_one_category_kernel(self, data_pack): """Get coco keypoint results.""" 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_kpt['keypoints'] for img_kpt in img_kpts]) key_points = _key_points.reshape(-1, self.ann_info['num_joints'] * 3) result = [{ 'image_id': img_kpt['image_id'], 'category_id': cat_id, 'keypoints': key_point.tolist(), 'score': float(img_kpt['score']), 'center': img_kpt['center'].tolist(), 'scale': img_kpt['scale'].tolist() } for img_kpt, key_point in zip(img_kpts, key_points)] cat_results.extend(result) return cat_results def _do_python_keypoint_eval(self, res_file): """Keypoint evaluation using COCOAPI.""" coco_det = self.coco.loadRes(res_file) coco_eval = COCOeval(self.coco, coco_det, 'keypoints', self.sigmas) 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 = list(zip(stats_names, coco_eval.stats)) return info_str def _sort_and_unique_bboxes(self, kpts, key='bbox_id'): """sort kpts and remove the repeated ones.""" for img_id, persons in kpts.items(): num = len(persons) kpts[img_id] = sorted(kpts[img_id], key=lambda x: x[key]) for i in range(num - 1, 0, -1): if kpts[img_id][i][key] == kpts[img_id][i - 1][key]: del kpts[img_id][i] return kpts
class BottomUpMhpDataset(BottomUpCocoDataset): """MHPv2.0 dataset for top-down pose estimation. `The Multi-Human Parsing project of Learning and Vision (LV) Group, National University of Singapore (NUS) is proposed to push the frontiers of fine-grained visual understanding of humans in crowd scene. <https://lv-mhp.github.io/>` The dataset loads raw features and apply specified transforms to return a dict containing the image tensors and other information. MHP keypoint indexes:: 0: "right ankle", 1: "right knee", 2: "right hip", 3: "left hip", 4: "left knee", 5: "left ankle", 6: "pelvis", 7: "thorax", 8: "upper neck", 9: "head top", 10: "right wrist", 11: "right elbow", 12: "right shoulder", 13: "left shoulder", 14: "left elbow", 15: "left wrist", Args: ann_file (str): Path to the annotation file. img_prefix (str): Path to a directory where images are held. Default: None. data_cfg (dict): config pipeline (list[dict | callable]): A sequence of data transforms. test_mode (bool): Store True when building test or validation dataset. Default: False. """ def __init__(self, ann_file, img_prefix, data_cfg, pipeline, test_mode=False): super(BottomUpCocoDataset, self).__init__(ann_file, img_prefix, data_cfg, pipeline, test_mode=test_mode) self.ann_info['flip_index'] = [ 5, 4, 3, 2, 1, 0, 6, 7, 8, 9, 15, 14, 13, 12, 11, 10 ] self.ann_info['use_different_joint_weights'] = False self.ann_info['joint_weights'] = np.array( [ 1.5, 1.2, 1., 1., 1.2, 1.5, 1., 1., 1., 1., 1.5, 1.2, 1., 1., 1.2, 1.5 ], dtype=np.float32).reshape((self.ann_info['num_joints'], 1)) # Adapted from COCO dataset self.sigmas = np.array([ .89, .83, 1.07, 1.07, .83, .89, .26, .26, .26, .26, .62, .72, 1.79, 1.79, .72, .62 ]) / 10.0 self.coco = COCO(ann_file) 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:]) self.img_ids = self.coco.getImgIds() if not test_mode: self.img_ids = [ img_id for img_id in self.img_ids if len(self.coco.getAnnIds(imgIds=img_id, iscrowd=None)) > 0 ] self.num_images = len(self.img_ids) self.id2name, self.name2id = self._get_mapping_id_name(self.coco.imgs) self.dataset_name = 'mhp' print(f'=> num_images: {self.num_images}') def _do_python_keypoint_eval(self, res_file): """Keypoint evaluation using COCOAPI.""" stats_names = [ 'AP', 'AP .5', 'AP .75', 'AP (M)', 'AP (L)', 'AR', 'AR .5', 'AR .75', 'AR (M)', 'AR (L)' ] with open(res_file, 'r') as file: res_json = json.load(file) if not res_json: info_str = list(zip(stats_names, [ 0, ] * len(stats_names))) return info_str coco_det = self.coco.loadRes(res_file) coco_eval = COCOeval(self.coco, coco_det, 'keypoints', self.sigmas, use_area=False) coco_eval.params.useSegm = None coco_eval.evaluate() coco_eval.accumulate() coco_eval.summarize() info_str = list(zip(stats_names, coco_eval.stats)) return info_str
def main(): parser = ArgumentParser() parser.add_argument('pose_config', help='Config file for pose network') parser.add_argument('pose_checkpoint', help='Checkpoint file') parser.add_argument('--img-root', type=str, default='', help='Image root') parser.add_argument('--json-file', type=str, default='', help='Json file containing image info.') parser.add_argument( '--camera-param-file', type=str, default=None, help='Camera parameter file for converting 3D pose predictions from ' ' the pixel space to camera space. If None, keypoints in pixel space' 'will be visualized') parser.add_argument( '--gt-joints-file', type=str, default=None, help='Optional arguement. Ground truth 3D keypoint parameter file. ' 'If None, gt keypoints will not be shown and keypoints in pixel ' 'space will be visualized.') parser.add_argument( '--rebase-keypoint-height', action='store_true', help='Rebase the predicted 3D pose so its lowest keypoint has a ' 'height of 0 (landing on the ground). This is useful for ' 'visualization when the model do not predict the global position ' 'of the 3D pose.') parser.add_argument( '--show-ground-truth', action='store_true', help='If True, show ground truth keypoint if it is available.') parser.add_argument('--show', action='store_true', default=False, help='whether to show img') parser.add_argument('--out-img-root', type=str, default=None, help='Root of the output visualization images. ' 'Default not saving the visualization images.') parser.add_argument('--device', default='cuda:0', help='Device for inference') parser.add_argument('--kpt-thr', type=float, default=0.3, help='Keypoint score threshold') parser.add_argument('--radius', type=int, default=4, help='Keypoint radius for visualization') parser.add_argument('--thickness', type=int, default=1, help='Link thickness for visualization') args = parser.parse_args() assert args.show or (args.out_img_root != '') coco = COCO(args.json_file) # build the pose model from a config file and a checkpoint file pose_model = init_pose_model(args.pose_config, args.pose_checkpoint, device=args.device.lower()) dataset = pose_model.cfg.data['test']['type'] # load camera parameters camera_params = None if args.camera_param_file is not None: camera_params = mmcv.load(args.camera_param_file) # load ground truth joints parameters gt_joint_params = None if args.gt_joints_file is not None: gt_joint_params = mmcv.load(args.gt_joints_file) # load hand bounding boxes det_results_list = [] for image_id, image in coco.imgs.items(): image_name = osp.join(args.img_root, image['file_name']) ann_ids = coco.getAnnIds(image_id) det_results = [] capture_key = str(image['capture']) camera_key = image['camera'] frame_idx = image['frame_idx'] for ann_id in ann_ids: ann = coco.anns[ann_id] if camera_params is not None: camera_param = { key: camera_params[capture_key][key][camera_key] for key in camera_params[capture_key].keys() } camera_param = _transform_interhand_camera_param(camera_param) else: camera_param = None if gt_joint_params is not None: joint_param = gt_joint_params[capture_key][str(frame_idx)] gt_joint = np.concatenate([ np.array(joint_param['world_coord']), np.array(joint_param['joint_valid']) ], axis=-1) else: gt_joint = None det_result = { 'image_name': image_name, 'bbox': ann['bbox'], # bbox format is 'xywh' 'camera_param': camera_param, 'keypoints_3d_gt': gt_joint } det_results.append(det_result) det_results_list.append(det_results) for i, det_results in enumerate( mmcv.track_iter_progress(det_results_list)): image_name = det_results[0]['image_name'] pose_results = inference_interhand_3d_model(pose_model, image_name, det_results, dataset=dataset) # Post processing pose_results_vis = [] for idx, res in enumerate(pose_results): keypoints_3d = res['keypoints_3d'] # normalize kpt score if keypoints_3d[:, 3].max() > 1: keypoints_3d[:, 3] /= 255 # get 2D keypoints in pixel space res['keypoints'] = keypoints_3d[:, [0, 1, 3]] # For model-predicted keypoints, channel 0 and 1 are coordinates # in pixel space, and channel 2 is the depth (in mm) relative # to root joints. # If both camera parameter and absolute depth of root joints are # provided, we can transform keypoint to camera space for better # visualization. camera_param = res['camera_param'] keypoints_3d_gt = res['keypoints_3d_gt'] if camera_param is not None and keypoints_3d_gt is not None: # build camera model camera = SimpleCamera(camera_param) # transform gt joints from world space to camera space keypoints_3d_gt[:, :3] = camera.world_to_camera( keypoints_3d_gt[:, :3]) # transform relative depth to absolute depth keypoints_3d[:21, 2] += keypoints_3d_gt[20, 2] keypoints_3d[21:, 2] += keypoints_3d_gt[41, 2] # transform keypoints from pixel space to camera space keypoints_3d[:, :3] = camera.pixel_to_camera( keypoints_3d[:, :3]) # rotate the keypoint to make z-axis correspondent to height # for better visualization vis_R = np.array([[1, 0, 0], [0, 0, -1], [0, 1, 0]]) keypoints_3d[:, :3] = keypoints_3d[:, :3] @ vis_R if keypoints_3d_gt is not None: keypoints_3d_gt[:, :3] = keypoints_3d_gt[:, :3] @ vis_R # rebase height (z-axis) if args.rebase_keypoint_height: valid = keypoints_3d[..., 3] > 0 keypoints_3d[..., 2] -= np.min(keypoints_3d[valid, 2], axis=-1, keepdims=True) res['keypoints_3d'] = keypoints_3d res['keypoints_3d_gt'] = keypoints_3d_gt # Add title instance_id = res.get('track_id', idx) res['title'] = f'Prediction ({instance_id})' pose_results_vis.append(res) # Add ground truth if args.show_ground_truth: if keypoints_3d_gt is None: print('Fail to show ground truth. Please make sure that' ' gt-joints-file is provided.') else: gt = res.copy() if args.rebase_keypoint_height: valid = keypoints_3d_gt[..., 3] > 0 keypoints_3d_gt[..., 2] -= np.min(keypoints_3d_gt[valid, 2], axis=-1, keepdims=True) gt['keypoints_3d'] = keypoints_3d_gt gt['title'] = f'Ground truth ({instance_id})' pose_results_vis.append(gt) # Visualization if args.out_img_root is None: out_file = None else: os.makedirs(args.out_img_root, exist_ok=True) out_file = osp.join(args.out_img_root, f'vis_{i}.jpg') vis_3d_pose_result( pose_model, result=pose_results_vis, img=det_results[0]['image_name'], out_file=out_file, dataset=dataset, show=args.show, kpt_score_thr=args.kpt_thr, radius=args.radius, thickness=args.thickness, )
def __init__(self, ann_file, img_prefix, data_cfg, pipeline, dataset_info=None, coco_style=True, test_mode=False): self.image_info = {} self.ann_info = {} self.ann_file = ann_file self.img_prefix = img_prefix self.pipeline = pipeline self.test_mode = test_mode # bottom-up self.base_size = data_cfg['base_size'] self.base_sigma = data_cfg['base_sigma'] self.int_sigma = False self.ann_info['image_size'] = np.array(data_cfg['image_size']) self.ann_info['heatmap_size'] = np.array(data_cfg['heatmap_size']) self.ann_info['num_joints'] = data_cfg['num_joints'] self.ann_info['num_scales'] = data_cfg['num_scales'] self.ann_info['scale_aware_sigma'] = data_cfg['scale_aware_sigma'] self.ann_info['inference_channel'] = data_cfg['inference_channel'] self.ann_info['dataset_channel'] = data_cfg['dataset_channel'] self.use_nms = data_cfg.get('use_nms', False) self.soft_nms = data_cfg.get('soft_nms', True) self.oks_thr = data_cfg.get('oks_thr', 0.9) if dataset_info is None: raise ValueError( 'Check https://github.com/open-mmlab/mmpose/pull/663 ' 'for details.') dataset_info = DatasetInfo(dataset_info) assert self.ann_info['num_joints'] == dataset_info.keypoint_num self.ann_info['flip_pairs'] = dataset_info.flip_pairs self.ann_info['flip_index'] = dataset_info.flip_index self.ann_info['upper_body_ids'] = dataset_info.upper_body_ids self.ann_info['lower_body_ids'] = dataset_info.lower_body_ids self.ann_info['joint_weights'] = dataset_info.joint_weights self.ann_info['skeleton'] = dataset_info.skeleton self.sigmas = dataset_info.sigmas self.dataset_name = dataset_info.dataset_name if coco_style: self.coco = COCO(ann_file) if 'categories' in self.coco.dataset: 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:]) self.img_ids = self.coco.getImgIds() if not test_mode: self.img_ids = [ img_id for img_id in self.img_ids if len(self.coco.getAnnIds(imgIds=img_id, iscrowd=None)) > 0 ] self.num_images = len(self.img_ids) self.id2name, self.name2id = self._get_mapping_id_name( self.coco.imgs) self.pipeline = Compose(self.pipeline)
class BottomUpCrowdPoseDataset(BottomUpCocoDataset): """CrowdPose dataset for bottom-up pose estimation. The dataset loads raw features and apply specified transforms to return a dict containing the image tensors and other information. CrowdPose keypoint indexes:: 0: 'left_shoulder', 1: 'right_shoulder', 2: 'left_elbow', 3: 'right_elbow', 4: 'left_wrist', 5: 'right_wrist', 6: 'left_hip', 7: 'right_hip', 8: 'left_knee', 9: 'right_knee', 10: 'left_ankle', 11: 'right_ankle', 12: 'top_head', 13: 'neck' Args: ann_file (str): Path to the annotation file. img_prefix (str): Path to a directory where images are held. Default: None. data_cfg (dict): config pipeline (list[dict | callable]): A sequence of data transforms. test_mode (bool): Store True when building test or validation dataset. Default: False. """ def __init__(self, ann_file, img_prefix, data_cfg, pipeline, test_mode=False): super(BottomUpCocoDataset, self).__init__(ann_file, img_prefix, data_cfg, pipeline, test_mode=test_mode) self.ann_info['flip_index'] = [ 1, 0, 3, 2, 5, 4, 7, 6, 9, 8, 11, 10, 12, 13 ] self.ann_info['use_different_joint_weights'] = False self.ann_info['joint_weights'] = np.array( [ 0.2, 0.2, 0.2, 1.3, 1.5, 0.2, 1.3, 1.5, 0.2, 0.2, 0.5, 0.2, 0.2, 0.5 ], dtype=np.float32).reshape((self.ann_info['num_joints'], 1)) self.sigmas = np.array([ .79, .79, .72, .72, .62, .62, 1.07, 1.07, .87, .87, .89, .89, .79, .79 ]) / 10.0 self.coco = COCO(ann_file) 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:]) self.img_ids = self.coco.getImgIds() if not test_mode: self.img_ids = [ img_id for img_id in self.img_ids if len(self.coco.getAnnIds(imgIds=img_id, iscrowd=None)) > 0 ] self.num_images = len(self.img_ids) self.id2name, self.name2id = self._get_mapping_id_name(self.coco.imgs) self.dataset_name = 'crowdpose' print(f'=> num_images: {self.num_images}') def _do_python_keypoint_eval(self, res_file): """Keypoint evaluation using COCOAPI.""" stats_names = [ 'AP', 'AP .5', 'AP .75', 'AR', 'AR .5', 'AR .75', 'AP(E)', 'AP(M)', 'AP(H)' ] with open(res_file, 'r') as file: res_json = json.load(file) if not res_json: info_str = list(zip(stats_names, [ 0, ] * len(stats_names))) return info_str coco_det = self.coco.loadRes(res_file) coco_eval = COCOeval(self.coco, coco_det, 'keypoints_crowd', self.sigmas, use_area=False) coco_eval.params.useSegm = None coco_eval.evaluate() coco_eval.accumulate() coco_eval.summarize() info_str = list(zip(stats_names, coco_eval.stats)) return info_str
class Kpt2dSviewRgbImgBottomUpDataset(Dataset, metaclass=ABCMeta): """Base class for bottom-up datasets. All datasets should subclass it. All subclasses should overwrite: Methods:`_get_single` Args: ann_file (str): Path to the annotation file. img_prefix (str): Path to a directory where images are held. Default: None. data_cfg (dict): config pipeline (list[dict | callable]): A sequence of data transforms. dataset_info (DatasetInfo): A class containing all dataset info. coco_style (bool): Whether the annotation json is coco-style. Default: True test_mode (bool): Store True when building test or validation dataset. Default: False. """ def __init__(self, ann_file, img_prefix, data_cfg, pipeline, dataset_info=None, coco_style=True, test_mode=False): self.image_info = {} self.ann_info = {} self.ann_file = ann_file self.img_prefix = img_prefix self.pipeline = pipeline self.test_mode = test_mode # bottom-up self.base_size = data_cfg['base_size'] self.base_sigma = data_cfg['base_sigma'] self.int_sigma = False self.ann_info['image_size'] = np.array(data_cfg['image_size']) self.ann_info['heatmap_size'] = np.array(data_cfg['heatmap_size']) self.ann_info['num_joints'] = data_cfg['num_joints'] self.ann_info['num_scales'] = data_cfg['num_scales'] self.ann_info['scale_aware_sigma'] = data_cfg['scale_aware_sigma'] self.ann_info['inference_channel'] = data_cfg['inference_channel'] self.ann_info['dataset_channel'] = data_cfg['dataset_channel'] self.use_nms = data_cfg.get('use_nms', False) self.soft_nms = data_cfg.get('soft_nms', True) self.oks_thr = data_cfg.get('oks_thr', 0.9) if dataset_info is None: raise ValueError( 'Check https://github.com/open-mmlab/mmpose/pull/663 ' 'for details.') dataset_info = DatasetInfo(dataset_info) assert self.ann_info['num_joints'] == dataset_info.keypoint_num self.ann_info['flip_pairs'] = dataset_info.flip_pairs self.ann_info['flip_index'] = dataset_info.flip_index self.ann_info['upper_body_ids'] = dataset_info.upper_body_ids self.ann_info['lower_body_ids'] = dataset_info.lower_body_ids self.ann_info['joint_weights'] = dataset_info.joint_weights self.ann_info['skeleton'] = dataset_info.skeleton self.sigmas = dataset_info.sigmas self.dataset_name = dataset_info.dataset_name if coco_style: self.coco = COCO(ann_file) if 'categories' in self.coco.dataset: 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:]) self.img_ids = self.coco.getImgIds() if not test_mode: self.img_ids = [ img_id for img_id in self.img_ids if len(self.coco.getAnnIds(imgIds=img_id, iscrowd=None)) > 0 ] self.num_images = len(self.img_ids) self.id2name, self.name2id = self._get_mapping_id_name( self.coco.imgs) self.pipeline = Compose(self.pipeline) @staticmethod def _get_mapping_id_name(imgs): """ Args: imgs (dict): dict of image info. Returns: tuple: Image name & id mapping dicts. - id2name (dict): Mapping image id to name. - name2id (dict): Mapping image name to id. """ id2name = {} name2id = {} for image_id, image in imgs.items(): file_name = image['file_name'] id2name[image_id] = file_name name2id[file_name] = image_id return id2name, name2id def _get_mask(self, anno, idx): """Get ignore masks to mask out losses.""" coco = self.coco img_info = coco.loadImgs(self.img_ids[idx])[0] m = np.zeros((img_info['height'], img_info['width']), dtype=np.float32) for obj in anno: if 'segmentation' in obj: if obj['iscrowd']: rle = xtcocotools.mask.frPyObjects(obj['segmentation'], img_info['height'], img_info['width']) m += xtcocotools.mask.decode(rle) elif obj['num_keypoints'] == 0: rles = xtcocotools.mask.frPyObjects( obj['segmentation'], img_info['height'], img_info['width']) for rle in rles: m += xtcocotools.mask.decode(rle) return m < 0.5 @abstractmethod def _get_single(self, idx): """Get anno for a single image.""" raise NotImplementedError @abstractmethod def evaluate(self, results, *args, **kwargs): """Evaluate keypoint results.""" def prepare_train_img(self, idx): """Prepare image for training given the index.""" results = copy.deepcopy(self._get_single(idx)) results['ann_info'] = self.ann_info return self.pipeline(results) def prepare_test_img(self, idx): """Prepare image for testing given the index.""" results = copy.deepcopy(self._get_single(idx)) results['ann_info'] = self.ann_info return self.pipeline(results) def __len__(self): """Get dataset length.""" return len(self.img_ids) def __getitem__(self, idx): """Get the sample for either training or testing given index.""" if self.test_mode: return self.prepare_test_img(idx) return self.prepare_train_img(idx)
def main(): """Visualize the demo images. Require the json_file containing boxes. """ parser = ArgumentParser() parser.add_argument('pose_config', help='Config file for detection') parser.add_argument('pose_checkpoint', help='Checkpoint file') parser.add_argument('--img-root', type=str, default='', help='Image root') parser.add_argument('--json-file', type=str, default='', help='Json file containing image info.') parser.add_argument('--show', action='store_true', default=False, help='whether to show img') parser.add_argument('--out-img-root', type=str, default='', help='Root of the output img file. ' 'Default not saving the visualization images.') parser.add_argument('--device', default='cuda:0', help='Device used for inference') parser.add_argument('--kpt-thr', type=float, default=0.3, help='Keypoint score threshold') args = parser.parse_args() assert args.show or (args.out_img_root != '') coco = COCO(args.json_file) # build the pose model from a config file and a checkpoint file pose_model = init_pose_model(args.pose_config, args.pose_checkpoint, device=args.device) dataset = pose_model.cfg.data['test']['type'] img_keys = list(coco.imgs.keys()) # optional return_heatmap = False # e.g. use ('backbone', ) to return backbone feature output_layer_names = None # process each image for i in range(len(img_keys)): # get bounding box annotations image_id = img_keys[i] image = coco.loadImgs(image_id)[0] image_name = os.path.join(args.img_root, image['file_name']) ann_ids = coco.getAnnIds(image_id) # make person bounding boxes person_bboxes = [] for ann_id in ann_ids: ann = coco.anns[ann_id] # bbox format is 'xywh' bbox = ann['bbox'] person_bboxes.append(bbox) # test a single image, with a list of bboxes pose_results, returned_outputs = inference_top_down_pose_model( pose_model, image_name, person_bboxes, bbox_thr=None, format='xywh', dataset=dataset, return_heatmap=return_heatmap, outputs=output_layer_names) if args.out_img_root == '': out_file = None else: os.makedirs(args.out_img_root, exist_ok=True) out_file = os.path.join(args.out_img_root, f'vis_{i}.jpg') vis_pose_result(pose_model, image_name, pose_results, dataset=dataset, kpt_score_thr=args.kpt_thr, show=args.show, out_file=out_file)
def main(): """Visualize the demo images. Require the json_file containing boxes. """ parser = ArgumentParser() parser.add_argument('pose_config', help='Config file for detection') parser.add_argument('pose_checkpoint', help='Checkpoint file') parser.add_argument('--img-root', type=str, default='', help='Image root') parser.add_argument('--json-file', type=str, default='', help='Json file containing image info.') parser.add_argument('--show', action='store_true', default=False, help='whether to show img') parser.add_argument('--out-img-root', type=str, default='', help='Root of the output img file. ' 'Default not saving the visualization images.') parser.add_argument('--device', default='cuda:0', help='Device used for inference') parser.add_argument('--kpt-thr', type=float, default=0.3, help='Keypoint score threshold') args = parser.parse_args() assert args.show or (args.out_img_root != '') skeleton = [[16, 14], [14, 12], [17, 15], [15, 13], [12, 13], [6, 12], [7, 13], [6, 7], [6, 8], [7, 9], [8, 10], [9, 11], [2, 3], [1, 2], [1, 3], [2, 4], [3, 5], [4, 6], [5, 7]] coco = COCO(args.json_file) # build the pose model from a config file and a checkpoint file pose_model = init_pose_model(args.pose_config, args.pose_checkpoint, device=args.device) img_keys = list(coco.imgs.keys()) # process each image for i in range(len(img_keys)): # get bounding box annotations image_id = img_keys[i] image = coco.loadImgs(image_id)[0] image_name = os.path.join(args.img_root, image['file_name']) ann_ids = coco.getAnnIds(image_id) # make person bounding boxes person_bboxes = [] for ann_id in ann_ids: ann = coco.anns[ann_id] # bbox format is 'xywh' bbox = ann['bbox'] person_bboxes.append(bbox) # test a single image, with a list of bboxes. pose_results = inference_top_down_pose_model(pose_model, image_name, person_bboxes, format='xywh') if args.out_img_root == '': out_file = None else: os.makedirs(args.out_img_root, exist_ok=True) out_file = os.path.join(args.out_img_root, f'vis_{i}.jpg') # show the results vis_pose_result(pose_model, image_name, pose_results, skeleton=skeleton, kpt_score_thr=args.kpt_thr, show=args.show, out_file=out_file)
class Kpt2dSviewRgbImgTopDownDataset(Dataset, metaclass=ABCMeta): """Base class for keypoint 2D top-down pose estimation with single-view RGB image as the input. All fashion datasets should subclass it. All subclasses should overwrite: Methods:`_get_db`, 'evaluate' Args: ann_file (str): Path to the annotation file. img_prefix (str): Path to a directory where images are held. Default: None. data_cfg (dict): config pipeline (list[dict | callable]): A sequence of data transforms. dataset_info (DatasetInfo): A class containing all dataset info. coco_style (bool): Whether the annotation json is coco-style. Default: True test_mode (bool): Store True when building test or validation dataset. Default: False. """ def __init__(self, ann_file, img_prefix, data_cfg, pipeline, dataset_info=None, coco_style=True, test_mode=False): self.image_info = {} self.ann_info = {} self.ann_file = ann_file self.img_prefix = img_prefix self.pipeline = pipeline self.test_mode = test_mode self.ann_info['image_size'] = np.array(data_cfg['image_size']) self.ann_info['heatmap_size'] = np.array(data_cfg['heatmap_size']) self.ann_info['num_joints'] = data_cfg['num_joints'] self.ann_info['inference_channel'] = data_cfg['inference_channel'] self.ann_info['num_output_channels'] = data_cfg['num_output_channels'] self.ann_info['dataset_channel'] = data_cfg['dataset_channel'] self.ann_info['use_different_joint_weights'] = data_cfg.get( 'use_different_joint_weights', False) if dataset_info is None: raise ValueError( 'Check https://github.com/open-mmlab/mmpose/pull/663 ' 'for details.') dataset_info = DatasetInfo(dataset_info) assert self.ann_info['num_joints'] == dataset_info.keypoint_num self.ann_info['flip_pairs'] = dataset_info.flip_pairs self.ann_info['flip_index'] = dataset_info.flip_index self.ann_info['upper_body_ids'] = dataset_info.upper_body_ids self.ann_info['lower_body_ids'] = dataset_info.lower_body_ids self.ann_info['joint_weights'] = dataset_info.joint_weights self.ann_info['skeleton'] = dataset_info.skeleton self.sigmas = dataset_info.sigmas self.dataset_name = dataset_info.dataset_name if coco_style: self.coco = COCO(ann_file) if 'categories' in self.coco.dataset: 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:]) self.img_ids = self.coco.getImgIds() self.num_images = len(self.img_ids) self.id2name, self.name2id = self._get_mapping_id_name( self.coco.imgs) self.db = [] self.pipeline = Compose(self.pipeline) @staticmethod def _get_mapping_id_name(imgs): """ Args: imgs (dict): dict of image info. Returns: tuple: Image name & id mapping dicts. - id2name (dict): Mapping image id to name. - name2id (dict): Mapping image name to id. """ id2name = {} name2id = {} for image_id, image in imgs.items(): file_name = image['file_name'] id2name[image_id] = file_name name2id[file_name] = image_id return id2name, name2id def _xywh2cs(self, x, y, w, h, padding=1.25): """This encodes bbox(x,y,w,h) into (center, scale) Args: x, y, w, h (float): left, top, width and height padding (float): bounding box padding factor Returns: center (np.ndarray[float32](2,)): center of the bbox (x, y). scale (np.ndarray[float32](2,)): scale of the bbox w & h. """ warnings.warn( 'The ``_xywh2cs`` method will be deprecated and removed from ' f'{self.__class__.__name__} in the future. Please use data ' 'transforms ``TopDownGetBboxCenterScale`` and ' '``TopDownRandomShiftBboxCenter`` in the pipeline instead.', DeprecationWarning) aspect_ratio = self.ann_info['image_size'][0] / self.ann_info[ 'image_size'][1] center = np.array([x + w * 0.5, y + h * 0.5], dtype=np.float32) if (not self.test_mode) and np.random.rand() < 0.3: center += 0.4 * (np.random.rand(2) - 0.5) * [w, h] if w > aspect_ratio * h: h = w * 1.0 / aspect_ratio elif w < aspect_ratio * h: w = h * aspect_ratio # pixel std is 200.0 scale = np.array([w / 200.0, h / 200.0], dtype=np.float32) # padding to include proper amount of context scale = scale * padding return center, scale def _get_normalize_factor(self, gts, *args, **kwargs): """Get the normalize factor. generally inter-ocular distance measured as the Euclidean distance between the outer corners of the eyes is used. This function should be overrode, to measure NME. Args: gts (np.ndarray[N, K, 2]): Groundtruth keypoint location. Returns: np.ndarray[N, 2]: normalized factor """ return np.ones([gts.shape[0], 2], dtype=np.float32) @abstractmethod def _get_db(self): """Load dataset.""" raise NotImplementedError @abstractmethod def evaluate(self, results, *args, **kwargs): """Evaluate keypoint results.""" @staticmethod def _write_keypoint_results(keypoints, res_file): """Write results into a json file.""" with open(res_file, 'w') as f: json.dump(keypoints, f, sort_keys=True, indent=4) def _report_metric(self, res_file, metrics, pck_thr=0.2, pckh_thr=0.7, auc_nor=30): """Keypoint evaluation. Args: res_file (str): Json file stored prediction results. metrics (str | list[str]): Metric to be performed. Options: 'PCK', 'PCKh', 'AUC', 'EPE', 'NME'. pck_thr (float): PCK threshold, default as 0.2. pckh_thr (float): PCKh threshold, default as 0.7. auc_nor (float): AUC normalization factor, default as 30 pixel. Returns: List: Evaluation results for evaluation metric. """ info_str = [] with open(res_file, 'r') as fin: preds = json.load(fin) assert len(preds) == len(self.db) outputs = [] gts = [] masks = [] box_sizes = [] threshold_bbox = [] threshold_head_box = [] for pred, item in zip(preds, self.db): outputs.append(np.array(pred['keypoints'])[:, :-1]) gts.append(np.array(item['joints_3d'])[:, :-1]) masks.append((np.array(item['joints_3d_visible'])[:, 0]) > 0) if 'PCK' in metrics: bbox = np.array(item['bbox']) bbox_thr = np.max(bbox[2:]) threshold_bbox.append(np.array([bbox_thr, bbox_thr])) if 'PCKh' in metrics: head_box_thr = item['head_size'] threshold_head_box.append( np.array([head_box_thr, head_box_thr])) box_sizes.append(item.get('box_size', 1)) outputs = np.array(outputs) gts = np.array(gts) masks = np.array(masks) threshold_bbox = np.array(threshold_bbox) threshold_head_box = np.array(threshold_head_box) box_sizes = np.array(box_sizes).reshape([-1, 1]) if 'PCK' in metrics: _, pck, _ = keypoint_pck_accuracy(outputs, gts, masks, pck_thr, threshold_bbox) info_str.append(('PCK', pck)) if 'PCKh' in metrics: _, pckh, _ = keypoint_pck_accuracy(outputs, gts, masks, pckh_thr, threshold_head_box) info_str.append(('PCKh', pckh)) if 'AUC' in metrics: info_str.append(('AUC', keypoint_auc(outputs, gts, masks, auc_nor))) if 'EPE' in metrics: info_str.append(('EPE', keypoint_epe(outputs, gts, masks))) if 'NME' in metrics: normalize_factor = self._get_normalize_factor(gts=gts, box_sizes=box_sizes) info_str.append( ('NME', keypoint_nme(outputs, gts, masks, normalize_factor))) return info_str def __len__(self): """Get the size of the dataset.""" return len(self.db) def __getitem__(self, idx): """Get the sample given index.""" results = copy.deepcopy(self.db[idx]) results['ann_info'] = self.ann_info return self.pipeline(results) def _sort_and_unique_bboxes(self, kpts, key='bbox_id'): """sort kpts and remove the repeated ones.""" kpts = sorted(kpts, key=lambda x: x[key]) num = len(kpts) for i in range(num - 1, 0, -1): if kpts[i][key] == kpts[i - 1][key]: del kpts[i] return kpts
class BottomUpAicDataset(BottomUpCocoDataset): """Aic dataset for bottom-up pose estimation. `AI Challenger : A Large-scale Dataset for Going Deeper in Image Understanding <https://arxiv.org/abs/1711.06475>`__ The dataset loads raw features and apply specified transforms to return a dict containing the image tensors and other information. AIC keypoint indexes:: 0: "right_shoulder", 1: "right_elbow", 2: "right_wrist", 3: "left_shoulder", 4: "left_elbow", 5: "left_wrist", 6: "right_hip", 7: "right_knee", 8: "right_ankle", 9: "left_hip", 10: "left_knee", 11: "left_ankle", 12: "head_top", 13: "neck" Args: ann_file (str): Path to the annotation file. img_prefix (str): Path to a directory where images are held. Default: None. data_cfg (dict): config pipeline (list[dict | callable]): A sequence of data transforms. test_mode (bool): Store True when building test or validation dataset. Default: False. """ def __init__(self, ann_file, img_prefix, data_cfg, pipeline, test_mode=False): super(BottomUpCocoDataset, self).__init__(ann_file, img_prefix, data_cfg, pipeline, test_mode=test_mode) self.ann_info['flip_index'] = [ 3, 4, 5, 0, 1, 2, 9, 10, 11, 6, 7, 8, 12, 13 ] self.ann_info['use_different_joint_weights'] = False self.ann_info['joint_weights'] = np.array( [1., 1.2, 1.5, 1., 1.2, 1.5, 1., 1.2, 1.5, 1., 1.2, 1.5, 1., 1.], dtype=np.float32).reshape((self.ann_info['num_joints'], 1)) self.sigmas = np.array([ 0.01388152, 0.01515228, 0.01057665, 0.01417709, 0.01497891, 0.01402144, 0.03909642, 0.03686941, 0.01981803, 0.03843971, 0.03412318, 0.02415081, 0.01291456, 0.01236173 ]) self.coco = COCO(ann_file) 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:]) self.img_ids = self.coco.getImgIds() if not test_mode: self.img_ids = [ img_id for img_id in self.img_ids if len(self.coco.getAnnIds(imgIds=img_id, iscrowd=None)) > 0 ] self.num_images = len(self.img_ids) self.id2name, self.name2id = self._get_mapping_id_name(self.coco.imgs) self.dataset_name = 'aic' print(f'=> num_images: {self.num_images}') def _do_python_keypoint_eval(self, res_file): """Keypoint evaluation using COCOAPI.""" stats_names = [ 'AP', 'AP .5', 'AP .75', 'AP (M)', 'AP (L)', 'AR', 'AR .5', 'AR .75', 'AR (M)', 'AR (L)' ] with open(res_file, 'r') as file: res_json = json.load(file) if not res_json: info_str = list(zip(stats_names, [ 0, ] * len(stats_names))) return info_str coco_det = self.coco.loadRes(res_file) coco_eval = COCOeval(self.coco, coco_det, 'keypoints', self.sigmas, use_area=False) coco_eval.params.useSegm = None coco_eval.evaluate() coco_eval.accumulate() coco_eval.summarize() info_str = list(zip(stats_names, coco_eval.stats)) return info_str