mmpose-0.29.0/tests/test_pipelines/test_shared_transform.py

# 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)