fhw
1 year ago
5 changed files with 425 additions and 12 deletions
@ -0,0 +1,407 @@ |
|||||
|
# Copyright (c) OpenMMLab. All rights reserved. |
||||
|
import os.path as osp |
||||
|
import tempfile |
||||
|
import warnings |
||||
|
from collections import OrderedDict, defaultdict |
||||
|
|
||||
|
import json_tricks as json |
||||
|
import numpy as np |
||||
|
from mmcv import Config, deprecated_api_warning |
||||
|
from xtcocotools.cocoeval import COCOeval |
||||
|
|
||||
|
from ....core.post_processing import oks_nms, soft_oks_nms |
||||
|
from ...builder import DATASETS |
||||
|
from ..base import Kpt2dSviewRgbImgTopDownDataset |
||||
|
|
||||
|
|
||||
|
@DATASETS.register_module() |
||||
|
class TopDownCocoDatasetSelf(Kpt2dSviewRgbImgTopDownDataset): |
||||
|
"""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. |
||||
|
dataset_info (DatasetInfo): A class containing all dataset info. |
||||
|
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, |
||||
|
test_mode=False): |
||||
|
|
||||
|
if dataset_info is None: |
||||
|
warnings.warn( |
||||
|
'dataset_info is missing. ' |
||||
|
'Check https://github.com/open-mmlab/mmpose/pull/663 ' |
||||
|
'for details.', DeprecationWarning) |
||||
|
cfg = Config.fromfile('configs/_base_/datasets/coco.py') |
||||
|
dataset_info = cfg._cfg_dict['dataset_info'] |
||||
|
|
||||
|
super().__init__( |
||||
|
ann_file, |
||||
|
img_prefix, |
||||
|
data_cfg, |
||||
|
pipeline, |
||||
|
dataset_info=dataset_info, |
||||
|
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) |
||||
|
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['sam_image_size'] = np.array(data_cfg['sam_image_size']) |
||||
|
|
||||
|
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.""" |
||||
|
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 = osp.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 _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 |
||||
|
|
||||
|
image_file = osp.join(self.img_prefix, |
||||
|
self.id2name[det_res['image_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 |
||||
|
|
||||
|
@deprecated_api_warning(name_dict=dict(outputs='results')) |
||||
|
def evaluate(self, results, res_folder=None, 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: |
||||
|
results (list[dict]): Testing results containing the following |
||||
|
items: |
||||
|
|
||||
|
- 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, optional): The folder to save the testing |
||||
|
results. If not specified, a temp folder will be created. |
||||
|
Default: None. |
||||
|
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') |
||||
|
|
||||
|
if res_folder is not None: |
||||
|
tmp_folder = None |
||||
|
res_file = osp.join(res_folder, 'result_keypoints.json') |
||||
|
else: |
||||
|
tmp_folder = tempfile.TemporaryDirectory() |
||||
|
res_file = osp.join(tmp_folder.name, 'result_keypoints.json') |
||||
|
|
||||
|
kpts = defaultdict(list) |
||||
|
|
||||
|
for result in results: |
||||
|
preds = result['preds'] |
||||
|
boxes = result['boxes'] |
||||
|
image_paths = result['image_paths'] |
||||
|
bbox_ids = result['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(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) |
||||
|
|
||||
|
if tmp_folder is not None: |
||||
|
tmp_folder.cleanup() |
||||
|
|
||||
|
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 |
||||
Loading…
Reference in new issue