def test_mpi3d_val_subset(mpi3d_data_dir):
data_specs = DataSpecs(
ImageSpecs(256,
mean=ImageSpecs.IMAGENET_MEAN,
stddev=ImageSpecs.IMAGENET_STDDEV),
JointsSpecs(CanonicalSkeletonDesc, n_dims=3),
)
dataset = MpiInf3dDataset(os.path.join(mpi3d_data_dir, 'val'), data_specs)
assert len(dataset) == 18561
example = dataset[0]
assert example['input'].shape == (3, 256, 256)
def get_dataset(dataset_name, data_specs=None, use_aug=False) -> PoseDataset:
"""Get a dataset instance by name.
Args:
dataset_name (str): The name of the dataset (eg. mpi3d-train)
data_specs (DataSpecs): The data specs for examples
use_aug (bool): If true, use data augmentation
Returns:
The pose dataset instance.
"""
# MPI-INF-3DHP
mpi3d_match = re.fullmatch('mpi3d-(train|val|test|test-uncorrected)',
dataset_name)
if mpi3d_match:
subset = mpi3d_match[1]
return MpiInf3dDataset(path.join(Base_Data_Dir, 'mpi3d', subset),
data_specs=data_specs,
use_aug=(use_aug
and not subset.startswith('test')))
if dataset_name == 'mpi3d-trainval':
return MixedPoseDataset([
get_dataset('mpi3d-train', data_specs, use_aug),
get_dataset('mpi3d-val', data_specs, use_aug),
],
balanced_sampling=False)
# Human3.6M
h36m_match = re.match('h36m-(trainval|test)', dataset_name)
if h36m_match:
subset = h36m_match[1]
return H36MDataset(path.join(Base_Data_Dir, 'h36m'),
data_specs=data_specs,
subset=subset,
use_aug=(use_aug and subset != 'test'))
# MPII Human Pose (2D)
mpii_match = re.match('mpii-(train|val|trainval|test)', dataset_name)
if mpii_match:
subset = mpii_match[1]
return MpiiDataset(path.join(Base_Data_Dir, 'mpii'),
data_specs=data_specs,
subset=subset,
use_aug=(use_aug and subset != 'test'))
raise Exception('unrecognised dataset: {}'.format(dataset_name))
def test_mpi3d_example_data(mpi3d_data_dir):
data_specs = DataSpecs(
ImageSpecs(256,
mean=ImageSpecs.IMAGENET_MEAN,
stddev=ImageSpecs.IMAGENET_STDDEV),
JointsSpecs(CanonicalSkeletonDesc, n_dims=3),
)
dataset = MpiInf3dDataset(os.path.join(mpi3d_data_dir, 'val'), data_specs)
assert MpiInf3dDataset.preserve_root_joint_at_univ_scale == False
example = dataset[0]
image = example['input']
assert float(image.min()) == pytest.approx(-2.117904, rel=0, abs=1e-2)
assert float(image.max()) == pytest.approx(2.428571, rel=0, abs=1e-2)
joints = example['target'][..., :3]
assert_allclose(joints[0], [-0.025768, -0.649297, -0.039933],
rtol=0,
atol=1e-4)
def test_to_canonical_skeleton(skeleton_mpi3d_univ, skeleton_canonical_univ):
actual = MpiInf3dDataset._mpi_inf_3dhp_to_canonical_skeleton(
skeleton_mpi3d_univ)
assert_allclose(actual, skeleton_canonical_univ)