show / mmpose-0.29.0 /tests /test_pipelines /test_shared_transform.py
camenduru's picture
thanks to show ❤
3bbb319
raw
history blame
6.71 kB
# Copyright (c) OpenMMLab. All rights reserved.
import os.path as osp
import numpy as np
import pytest
from mmcv import bgr2rgb, build_from_cfg
from mmpose.datasets import PIPELINES
from mmpose.datasets.pipelines import Compose
def check_keys_equal(result_keys, target_keys):
"""Check if all elements in target_keys is in result_keys."""
return set(target_keys) == set(result_keys)
def check_keys_contain(result_keys, target_keys):
"""Check if elements in target_keys is in result_keys."""
return set(target_keys).issubset(set(result_keys))
def test_compose():
with pytest.raises(TypeError):
# transform must be callable or a dict
Compose('LoadImageFromFile')
target_keys = ['img', 'img_rename', 'img_metas']
# test Compose given a data pipeline
img = np.random.randn(256, 256, 3)
results = dict(img=img, img_file='test_image.png')
test_pipeline = [
dict(
type='Collect',
keys=['img', ('img', 'img_rename')],
meta_keys=['img_file'])
]
compose = Compose(test_pipeline)
compose_results = compose(results)
assert check_keys_equal(compose_results.keys(), target_keys)
assert check_keys_equal(compose_results['img_metas'].data.keys(),
['img_file'])
# test Compose when forward data is None
results = None
class ExamplePipeline:
def __call__(self, results):
return None
nonePipeline = ExamplePipeline()
test_pipeline = [nonePipeline]
compose = Compose(test_pipeline)
compose_results = compose(results)
assert compose_results is None
assert repr(compose) == compose.__class__.__name__ + \
f'(\n {nonePipeline}\n)'
def test_load_image_from_file():
# Define simple pipeline
load = dict(type='LoadImageFromFile')
load = build_from_cfg(load, PIPELINES)
data_prefix = 'tests/data/coco/'
image_file = osp.join(data_prefix, '00000000078.jpg')
results = dict(image_file=image_file)
# load an image that doesn't exist
with pytest.raises(FileNotFoundError):
results = load(results)
# mormal loading
image_file = osp.join(data_prefix, '000000000785.jpg')
results = dict(image_file=image_file)
results = load(results)
assert results['img'].shape == (425, 640, 3)
# load a single image from a list
image_file = [osp.join(data_prefix, '000000000785.jpg')]
results = dict(image_file=image_file)
results = load(results)
assert len(results['img']) == 1
# test loading multi images from a list
image_file = [
osp.join(data_prefix, '000000000785.jpg'),
osp.join(data_prefix, '00000004008.jpg'),
]
results = dict(image_file=image_file)
with pytest.raises(FileNotFoundError):
results = load(results)
image_file = [
osp.join(data_prefix, '000000000785.jpg'),
osp.join(data_prefix, '000000040083.jpg'),
]
results = dict(image_file=image_file)
results = load(results)
assert len(results['img']) == 2
# manually set image outside the pipeline
img = np.random.randint(0, 255, (32, 32, 3), dtype=np.uint8)
results = load(dict(img=img))
np.testing.assert_equal(results['img'], bgr2rgb(img))
imgs = np.random.randint(0, 255, (2, 32, 32, 3), dtype=np.uint8)
desired = np.concatenate([bgr2rgb(img) for img in imgs], axis=0)
results = load(dict(img=imgs))
np.testing.assert_equal(results['img'], desired)
# neither 'image_file' or valid 'img' is given
results = dict()
with pytest.raises(KeyError):
_ = load(results)
results = dict(img=np.random.randint(0, 255, (32, 32), dtype=np.uint8))
with pytest.raises(ValueError):
_ = load(results)
def test_albu_transform():
data_prefix = 'tests/data/coco/'
results = dict(image_file=osp.join(data_prefix, '000000000785.jpg'))
# Define simple pipeline
load = dict(type='LoadImageFromFile')
load = build_from_cfg(load, PIPELINES)
albu_transform = dict(
type='Albumentation',
transforms=[
dict(type='RandomBrightnessContrast', p=0.2),
dict(type='ToFloat')
])
albu_transform = build_from_cfg(albu_transform, PIPELINES)
# Execute transforms
results = load(results)
results = albu_transform(results)
assert results['img'].dtype == np.float32
def test_photometric_distortion_transform():
data_prefix = 'tests/data/coco/'
results = dict(image_file=osp.join(data_prefix, '000000000785.jpg'))
# Define simple pipeline
load = dict(type='LoadImageFromFile')
load = build_from_cfg(load, PIPELINES)
photo_transform = dict(type='PhotometricDistortion')
photo_transform = build_from_cfg(photo_transform, PIPELINES)
# Execute transforms
results = load(results)
results = photo_transform(results)
assert results['img'].dtype == np.uint8
def test_multitask_gather():
ann_info = dict(
image_size=np.array([256, 256]),
heatmap_size=np.array([64, 64]),
num_joints=17,
joint_weights=np.ones((17, 1), dtype=np.float32),
use_different_joint_weights=False)
results = dict(
joints_3d=np.zeros([17, 3]),
joints_3d_visible=np.ones([17, 3]),
ann_info=ann_info)
pipeline_list = [[dict(type='TopDownGenerateTarget', sigma=2)],
[dict(type='TopDownGenerateTargetRegression')]]
pipeline = dict(
type='MultitaskGatherTarget',
pipeline_list=pipeline_list,
pipeline_indices=[0, 1, 0],
)
pipeline = build_from_cfg(pipeline, PIPELINES)
results = pipeline(results)
target = results['target']
target_weight = results['target_weight']
assert isinstance(target, list)
assert isinstance(target_weight, list)
assert target[0].shape == (17, 64, 64)
assert target_weight[0].shape == (17, 1)
assert target[1].shape == (17, 2)
assert target_weight[1].shape == (17, 2)
assert target[2].shape == (17, 64, 64)
assert target_weight[2].shape == (17, 1)
def test_rename_keys():
results = dict(
joints_3d=np.ones([17, 3]), joints_3d_visible=np.ones([17, 3]))
pipeline = dict(
type='RenameKeys',
key_pairs=[('joints_3d', 'target'),
('joints_3d_visible', 'target_weight')])
pipeline = build_from_cfg(pipeline, PIPELINES)
results = pipeline(results)
assert 'joints_3d' not in results
assert 'joints_3d_visible' not in results
assert 'target' in results
assert 'target_weight' in results
assert results['target'].shape == (17, 3)
assert results['target_weight'].shape == (17, 3)