You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
255 lines
9.2 KiB
255 lines
9.2 KiB
# Copyright (c) OpenMMLab. All rights reserved.
|
|
import copy
|
|
import os
|
|
|
|
import numpy as np
|
|
import torch
|
|
from numpy.testing import assert_array_almost_equal
|
|
|
|
from mmpose.datasets.pipelines import (Collect, IUVToTensor, LoadImageFromFile,
|
|
LoadIUVFromFile, MeshAffine,
|
|
MeshGetRandomScaleRotation,
|
|
MeshRandomChannelNoise, MeshRandomFlip,
|
|
NormalizeTensor, ToTensor)
|
|
|
|
|
|
def _check_keys_contain(result_keys, target_keys):
|
|
"""Check if all elements in target_keys is in result_keys."""
|
|
return set(target_keys).issubset(set(result_keys))
|
|
|
|
|
|
def _check_flip(origin_imgs, result_imgs):
|
|
"""Check if the origin_imgs are flipped correctly."""
|
|
h, w, c = origin_imgs.shape
|
|
for i in range(h):
|
|
for j in range(w):
|
|
for k in range(c):
|
|
if result_imgs[i, j, k] != origin_imgs[i, w - 1 - j, k]:
|
|
return False
|
|
return True
|
|
|
|
|
|
def _check_rot90(origin_imgs, result_imgs):
|
|
if origin_imgs.shape[0] == result_imgs.shape[1] and \
|
|
origin_imgs.shape[1] == result_imgs.shape[0]:
|
|
return True
|
|
else:
|
|
return False
|
|
|
|
|
|
def _check_normalize(origin_imgs, result_imgs, norm_cfg):
|
|
"""Check if the origin_imgs are normalized correctly into result_imgs in a
|
|
given norm_cfg."""
|
|
target_imgs = result_imgs.copy()
|
|
for i in range(3):
|
|
target_imgs[i] *= norm_cfg['std'][i]
|
|
target_imgs[i] += norm_cfg['mean'][i]
|
|
assert_array_almost_equal(origin_imgs, target_imgs, decimal=4)
|
|
|
|
|
|
def _box2cs(box, image_size):
|
|
x, y, w, h = box[:4]
|
|
|
|
aspect_ratio = 1. * image_size[0] / image_size[1]
|
|
center = np.zeros((2), dtype=np.float32)
|
|
center[0] = x + w * 0.5
|
|
center[1] = y + h * 0.5
|
|
|
|
if w > aspect_ratio * h:
|
|
h = w * 1.0 / aspect_ratio
|
|
elif w < aspect_ratio * h:
|
|
w = h * aspect_ratio
|
|
scale = np.array([w * 1.0 / 200.0, h * 1.0 / 200.0], dtype=np.float32)
|
|
scale = scale * 1.25
|
|
return center, scale
|
|
|
|
|
|
def _load_test_data():
|
|
data_cfg = dict(
|
|
image_size=[256, 256],
|
|
iuv_size=[64, 64],
|
|
num_joints=24,
|
|
use_IUV=True,
|
|
uv_type='BF')
|
|
ann_file = 'tests/data/h36m/test_h36m.npz'
|
|
img_prefix = 'tests/data/h36m'
|
|
index = 0
|
|
|
|
ann_info = dict(image_size=np.array(data_cfg['image_size']))
|
|
ann_info['iuv_size'] = np.array(data_cfg['iuv_size'])
|
|
ann_info['num_joints'] = data_cfg['num_joints']
|
|
ann_info['flip_pairs'] = [[0, 5], [1, 4], [2, 3], [6, 11], [7, 10], [8, 9],
|
|
[20, 21], [22, 23]]
|
|
ann_info['use_different_joint_weights'] = False
|
|
ann_info['joint_weights'] = \
|
|
np.ones(ann_info['num_joints'], dtype=np.float32
|
|
).reshape(ann_info['num_joints'], 1)
|
|
ann_info['uv_type'] = data_cfg['uv_type']
|
|
ann_info['use_IUV'] = data_cfg['use_IUV']
|
|
uv_type = ann_info['uv_type']
|
|
iuv_prefix = os.path.join(img_prefix, f'{uv_type}_IUV_gt')
|
|
|
|
ann_data = np.load(ann_file)
|
|
|
|
results = dict(ann_info=ann_info)
|
|
results['rotation'] = 0
|
|
results['image_file'] = os.path.join(img_prefix,
|
|
ann_data['imgname'][index])
|
|
scale = ann_data['scale'][index]
|
|
results['scale'] = np.array([scale, scale]).astype(np.float32)
|
|
results['center'] = ann_data['center'][index].astype(np.float32)
|
|
|
|
# Get gt 2D joints, if available
|
|
if 'part' in ann_data.keys():
|
|
keypoints = ann_data['part'][index].astype(np.float32)
|
|
results['joints_2d'] = keypoints[:, :2]
|
|
results['joints_2d_visible'] = keypoints[:, -1][:, np.newaxis]
|
|
else:
|
|
results['joints_2d'] = np.zeros((24, 2), dtype=np.float32)
|
|
results['joints_2d_visible'] = np.zeros((24, 1), dtype=np.float32)
|
|
|
|
# Get gt 3D joints, if available
|
|
if 'S' in ann_data.keys():
|
|
joints_3d = ann_data['S'][index].astype(np.float32)
|
|
results['joints_3d'] = joints_3d[:, :3]
|
|
results['joints_3d_visible'] = joints_3d[:, -1][:, np.newaxis]
|
|
else:
|
|
results['joints_3d'] = np.zeros((24, 3), dtype=np.float32)
|
|
results['joints_3d_visible'] = np.zeros((24, 1), dtype=np.float32)
|
|
|
|
# Get gt SMPL parameters, if available
|
|
if 'pose' in ann_data.keys() and 'shape' in ann_data.keys():
|
|
results['pose'] = ann_data['pose'][index].astype(np.float32)
|
|
results['beta'] = ann_data['shape'][index].astype(np.float32)
|
|
results['has_smpl'] = 1
|
|
else:
|
|
results['pose'] = np.zeros(72, dtype=np.float32)
|
|
results['beta'] = np.zeros(10, dtype=np.float32)
|
|
results['has_smpl'] = 0
|
|
|
|
# Get gender data, if available
|
|
if 'gender' in ann_data.keys():
|
|
gender = ann_data['gender'][index]
|
|
results['gender'] = 0 if str(gender) == 'm' else 1
|
|
else:
|
|
results['gender'] = -1
|
|
|
|
# Get IUV image, if available
|
|
if 'iuv_names' in ann_data.keys():
|
|
results['iuv_file'] = os.path.join(iuv_prefix,
|
|
ann_data['iuv_names'][index])
|
|
results['has_iuv'] = results['has_smpl']
|
|
else:
|
|
results['iuv_file'] = ''
|
|
results['has_iuv'] = 0
|
|
|
|
return copy.deepcopy(results)
|
|
|
|
|
|
def test_mesh_pipeline():
|
|
# load data
|
|
results = _load_test_data()
|
|
|
|
# data_prefix = 'tests/data/coco/'
|
|
# ann_file = osp.join(data_prefix, 'test_coco.json')
|
|
# coco = COCO(ann_file)
|
|
#
|
|
# results = dict(image_file=osp.join(data_prefix, '000000000785.jpg'))
|
|
|
|
# test loading image
|
|
transform = LoadImageFromFile()
|
|
results = transform(copy.deepcopy(results))
|
|
assert results['img'].shape == (1002, 1000, 3)
|
|
|
|
# test loading densepose IUV image without GT iuv image
|
|
transform = LoadIUVFromFile()
|
|
results_no_iuv = copy.deepcopy(results)
|
|
results_no_iuv['has_iuv'] = 0
|
|
results_no_iuv = transform(results_no_iuv)
|
|
assert results_no_iuv['iuv'] is None
|
|
|
|
# test loading densepose IUV image
|
|
results = transform(results)
|
|
assert results['iuv'].shape == (1002, 1000, 3)
|
|
assert results['iuv'][:, :, 0].max() <= 1
|
|
|
|
# test flip
|
|
random_flip = MeshRandomFlip(flip_prob=1.)
|
|
results_flip = random_flip(copy.deepcopy(results))
|
|
assert _check_flip(results['img'], results_flip['img'])
|
|
flip_iuv = results_flip['iuv']
|
|
flip_iuv[:, :, 1] = 255 - flip_iuv[:, :, 1]
|
|
assert _check_flip(results['iuv'], flip_iuv)
|
|
results = results_flip
|
|
|
|
# test flip without IUV image
|
|
results_no_iuv = random_flip(copy.deepcopy(results_no_iuv))
|
|
assert results_no_iuv['iuv'] is None
|
|
|
|
# test random scale and rotation
|
|
random_scale_rotation = MeshGetRandomScaleRotation()
|
|
results = random_scale_rotation(results)
|
|
|
|
# test affine
|
|
affine_transform = MeshAffine()
|
|
results_affine = affine_transform(copy.deepcopy(results))
|
|
assert results_affine['img'].shape == (256, 256, 3)
|
|
assert results_affine['iuv'].shape == (64, 64, 3)
|
|
results = results_affine
|
|
|
|
# test affine without IUV image
|
|
results_no_iuv['rotation'] = 30
|
|
results_no_iuv = affine_transform(copy.deepcopy(results_no_iuv))
|
|
assert results_no_iuv['iuv'] is None
|
|
|
|
# test channel noise
|
|
random_noise = MeshRandomChannelNoise()
|
|
results_noise = random_noise(copy.deepcopy(results))
|
|
results = results_noise
|
|
|
|
# transfer image to tensor
|
|
to_tensor = ToTensor()
|
|
results_tensor = to_tensor(copy.deepcopy(results))
|
|
assert isinstance(results_tensor['img'], torch.Tensor)
|
|
assert results_tensor['img'].shape == torch.Size([3, 256, 256])
|
|
|
|
# transfer IUV image to tensor
|
|
iuv_to_tensor = IUVToTensor()
|
|
results_tensor = iuv_to_tensor(results_tensor)
|
|
assert isinstance(results_tensor['part_index'], torch.LongTensor)
|
|
assert results_tensor['part_index'].shape == torch.Size([1, 64, 64])
|
|
max_I = results_tensor['part_index'].max().item()
|
|
assert (max_I == 0 or max_I == 1)
|
|
assert isinstance(results_tensor['uv_coordinates'], torch.FloatTensor)
|
|
assert results_tensor['uv_coordinates'].shape == torch.Size([2, 64, 64])
|
|
|
|
# transfer IUV image to tensor without GT IUV image
|
|
results_no_iuv = iuv_to_tensor(results_no_iuv)
|
|
assert isinstance(results_no_iuv['part_index'], torch.LongTensor)
|
|
assert results_no_iuv['part_index'].shape == torch.Size([1, 64, 64])
|
|
max_I = results_no_iuv['part_index'].max().item()
|
|
assert (max_I == 0)
|
|
assert isinstance(results_no_iuv['uv_coordinates'], torch.FloatTensor)
|
|
assert results_no_iuv['uv_coordinates'].shape == torch.Size([2, 64, 64])
|
|
|
|
# test norm
|
|
norm_cfg = {}
|
|
norm_cfg['mean'] = [0.485, 0.456, 0.406]
|
|
norm_cfg['std'] = [0.229, 0.224, 0.225]
|
|
normalize = NormalizeTensor(mean=norm_cfg['mean'], std=norm_cfg['std'])
|
|
|
|
results_normalize = normalize(copy.deepcopy(results_tensor))
|
|
_check_normalize(results_tensor['img'].data.numpy(),
|
|
results_normalize['img'].data.numpy(), norm_cfg)
|
|
|
|
# test collect
|
|
collect = Collect(
|
|
keys=[
|
|
'img', 'joints_2d', 'joints_2d_visible', 'joints_3d',
|
|
'joints_3d_visible', 'pose', 'beta', 'part_index', 'uv_coordinates'
|
|
],
|
|
meta_keys=['image_file', 'center', 'scale', 'rotation', 'iuv_file'])
|
|
results_final = collect(results_normalize)
|
|
|
|
assert 'img_size' not in results_final['img_metas'].data
|
|
assert 'image_file' in results_final['img_metas'].data
|
|
|