|
|
|
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): |
|
|
|
Compose('LoadImageFromFile') |
|
|
|
target_keys = ['img', 'img_rename', 'img_metas'] |
|
|
|
|
|
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']) |
|
|
|
|
|
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(): |
|
|
|
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) |
|
|
|
|
|
with pytest.raises(FileNotFoundError): |
|
results = load(results) |
|
|
|
|
|
image_file = osp.join(data_prefix, '000000000785.jpg') |
|
results = dict(image_file=image_file) |
|
results = load(results) |
|
assert results['img'].shape == (425, 640, 3) |
|
|
|
|
|
image_file = [osp.join(data_prefix, '000000000785.jpg')] |
|
results = dict(image_file=image_file) |
|
results = load(results) |
|
assert len(results['img']) == 1 |
|
|
|
|
|
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 |
|
|
|
|
|
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) |
|
|
|
|
|
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')) |
|
|
|
|
|
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) |
|
|
|
|
|
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')) |
|
|
|
|
|
load = dict(type='LoadImageFromFile') |
|
load = build_from_cfg(load, PIPELINES) |
|
|
|
photo_transform = dict(type='PhotometricDistortion') |
|
photo_transform = build_from_cfg(photo_transform, PIPELINES) |
|
|
|
|
|
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) |
|
|