Source code for mmpose.datasets.datasets.body.mpii_dataset
# Copyright (c) OpenMMLab. All rights reserved.
import json
import os.path as osp
from typing import Callable, List, Optional, Sequence, Tuple, Union
import numpy as np
from mmengine.fileio import exists, get_local_path
from scipy.io import loadmat
from mmpose.registry import DATASETS
from mmpose.structures.bbox import bbox_cs2xyxy
from ..base import BaseCocoStyleDataset
[docs]@DATASETS.register_module()
class MpiiDataset(BaseCocoStyleDataset):
"""MPII Dataset for pose estimation.
"2D Human Pose Estimation: New Benchmark and State of the Art Analysis"
,CVPR'2014. More details can be found in the `paper
<http://human-pose.mpi-inf.mpg.de/contents/andriluka14cvpr.pdf>`__ .
MPII keypoints::
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): Annotation file path. Default: ''.
bbox_file (str, optional): Detection result file path. If
``bbox_file`` is set, detected bboxes loaded from this file will
be used instead of ground-truth bboxes. This setting is only for
evaluation, i.e., ignored when ``test_mode`` is ``False``.
Default: ``None``.
headbox_file (str, optional): The path of ``mpii_gt_val.mat`` which
provides the headboxes information used for ``PCKh``.
Default: ``None``.
data_mode (str): Specifies the mode of data samples: ``'topdown'`` or
``'bottomup'``. In ``'topdown'`` mode, each data sample contains
one instance; while in ``'bottomup'`` mode, each data sample
contains all instances in a image. Default: ``'topdown'``
metainfo (dict, optional): Meta information for dataset, such as class
information. Default: ``None``.
data_root (str, optional): The root directory for ``data_prefix`` and
``ann_file``. Default: ``None``.
data_prefix (dict, optional): Prefix for training data. Default:
``dict(img=None, ann=None)``.
filter_cfg (dict, optional): Config for filter data. Default: `None`.
indices (int or Sequence[int], optional): Support using first few
data in annotation file to facilitate training/testing on a smaller
dataset. Default: ``None`` which means using all ``data_infos``.
serialize_data (bool, optional): Whether to hold memory using
serialized objects, when enabled, data loader workers can use
shared RAM from master process instead of making a copy.
Default: ``True``.
pipeline (list, optional): Processing pipeline. Default: [].
test_mode (bool, optional): ``test_mode=True`` means in test phase.
Default: ``False``.
lazy_init (bool, optional): Whether to load annotation during
instantiation. In some cases, such as visualization, only the meta
information of the dataset is needed, which is not necessary to
load annotation file. ``Basedataset`` can skip load annotations to
save time by set ``lazy_init=False``. Default: ``False``.
max_refetch (int, optional): If ``Basedataset.prepare_data`` get a
None img. The maximum extra number of cycles to get a valid
image. Default: 1000.
"""
METAINFO: dict = dict(from_file='configs/_base_/datasets/mpii.py')
def __init__(self,
ann_file: str = '',
bbox_file: Optional[str] = None,
headbox_file: Optional[str] = None,
data_mode: str = 'topdown',
metainfo: Optional[dict] = None,
data_root: Optional[str] = None,
data_prefix: dict = dict(img=''),
filter_cfg: Optional[dict] = None,
indices: Optional[Union[int, Sequence[int]]] = None,
serialize_data: bool = True,
pipeline: List[Union[dict, Callable]] = [],
test_mode: bool = False,
lazy_init: bool = False,
max_refetch: int = 1000):
if headbox_file:
if data_mode != 'topdown':
raise ValueError(
f'{self.__class__.__name__} is set to {data_mode}: '
'mode, while "headbox_file" is only '
'supported in topdown mode.')
if not test_mode:
raise ValueError(
f'{self.__class__.__name__} has `test_mode==False` '
'while "headbox_file" is only '
'supported when `test_mode==True`.')
headbox_file_type = headbox_file[-3:]
allow_headbox_file_type = ['mat']
if headbox_file_type not in allow_headbox_file_type:
raise KeyError(
f'The head boxes file type {headbox_file_type} is not '
f'supported. Should be `mat` but got {headbox_file_type}.')
self.headbox_file = headbox_file
super().__init__(
ann_file=ann_file,
bbox_file=bbox_file,
data_mode=data_mode,
metainfo=metainfo,
data_root=data_root,
data_prefix=data_prefix,
filter_cfg=filter_cfg,
indices=indices,
serialize_data=serialize_data,
pipeline=pipeline,
test_mode=test_mode,
lazy_init=lazy_init,
max_refetch=max_refetch)
def _load_annotations(self) -> Tuple[List[dict], List[dict]]:
"""Load data from annotations in MPII format."""
assert exists(self.ann_file), (
f'Annotation file `{self.ann_file}` does not exist')
with get_local_path(self.ann_file) as local_path:
with open(local_path) as anno_file:
self.anns = json.load(anno_file)
if self.headbox_file:
assert exists(self.headbox_file), (
f'Headbox file `{self.headbox_file}` does not exist')
with get_local_path(self.headbox_file) as local_path:
self.headbox_dict = loadmat(local_path)
headboxes_src = np.transpose(self.headbox_dict['headboxes_src'],
[2, 0, 1])
SC_BIAS = 0.6
instance_list = []
image_list = []
used_img_ids = set()
ann_id = 0
# mpii bbox scales are normalized with factor 200.
pixel_std = 200.
for idx, ann in enumerate(self.anns):
center = np.array(ann['center'], dtype=np.float32)
scale = np.array([ann['scale'], ann['scale']],
dtype=np.float32) * pixel_std
# Adjust center/scale slightly to avoid cropping limbs
if center[0] != -1:
center[1] = center[1] + 15. / pixel_std * scale[1]
# MPII uses matlab format, index is 1-based,
# we should first convert to 0-based index
center = center - 1
# unify shape with coco datasets
center = center.reshape(1, -1)
scale = scale.reshape(1, -1)
bbox = bbox_cs2xyxy(center, scale)
# load keypoints in shape [1, K, 2] and keypoints_visible in [1, K]
keypoints = np.array(
ann['joints'], dtype=np.float32).reshape(1, -1, 2)
keypoints_visible = np.array(ann['joints_vis']).reshape(1, -1)
x1, y1, x2, y2 = np.split(bbox, axis=1, indices_or_sections=4)
area = np.clip((x2 - x1) * (y2 - y1) * 0.53, a_min=1.0, a_max=None)
area = area[..., 0].astype(np.float32)
category_id = ann.get('category_id', [1] * len(bbox))
instance_info = {
'id': ann_id,
'img_id': int(ann['image'].split('.')[0]),
'img_path': osp.join(self.data_prefix['img'], ann['image']),
'bbox_center': center,
'bbox_scale': scale,
'bbox': bbox,
'bbox_score': np.ones(1, dtype=np.float32),
'keypoints': keypoints,
'keypoints_visible': keypoints_visible,
'area': area,
'category_id': category_id,
}
if self.headbox_file:
# calculate the diagonal length of head box as norm_factor
headbox = headboxes_src[idx]
head_size = np.linalg.norm(headbox[1] - headbox[0], axis=0)
head_size *= SC_BIAS
instance_info['head_size'] = head_size.reshape(1, -1)
if instance_info['img_id'] not in used_img_ids:
used_img_ids.add(instance_info['img_id'])
image_list.append({
'img_id': instance_info['img_id'],
'img_path': instance_info['img_path'],
})
instance_list.append(instance_info)
ann_id = ann_id + 1
return instance_list, image_list