fhw
/
ViTPose
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity
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.
14 Commits
2 Branches
0 Tags
59 MiB
 
Tree: 98fb7f859f
Branches Tags
${ item.name }
Create tag ${ searchTerm }
Create branch ${ searchTerm }
from '98fb7f859f'
${ noResults }
ViTPose/tests/test_pipelines/test_bottom_up_pipelines.py

427 lines
17 KiB
Raw Blame History

# Copyright (c) OpenMMLab. All rights reserved.
import copy
import os.path as osp
import numpy as np
import pytest
import xtcocotools
from xtcocotools.coco import COCO
from mmpose.datasets.pipelines import (BottomUpGenerateHeatmapTarget,
BottomUpGeneratePAFTarget,
BottomUpGenerateTarget,
BottomUpGetImgSize,
BottomUpRandomAffine,
BottomUpRandomFlip, BottomUpResizeAlign,
LoadImageFromFile)
def _get_mask(coco, anno, img_id):
img_info = coco.loadImgs(img_id)[0]
m = np.zeros((img_info['height'], img_info['width']), dtype=np.float32)
for obj in anno:
if obj['iscrowd']:
rle = xtcocotools.mask.frPyObjects(obj['segmentation'],
img_info['height'],
img_info['width'])
m += xtcocotools.mask.decode(rle)
elif obj['num_keypoints'] == 0:
rles = xtcocotools.mask.frPyObjects(obj['segmentation'],
img_info['height'],
img_info['width'])
for rle in rles:
m += xtcocotools.mask.decode(rle)
return m < 0.5
def _get_joints(anno, ann_info, int_sigma):
num_people = len(anno)
if ann_info['scale_aware_sigma']:
joints = np.zeros((num_people, ann_info['num_joints'], 4),
dtype=np.float32)
else:
joints = np.zeros((num_people, ann_info['num_joints'], 3),
dtype=np.float32)
for i, obj in enumerate(anno):
joints[i, :ann_info['num_joints'], :3] = \
np.array(obj['keypoints']).reshape([-1, 3])
if ann_info['scale_aware_sigma']:
# get person box
box = obj['bbox']
size = max(box[2], box[3])
sigma = size / 256 * 2
if int_sigma:
sigma = int(np.ceil(sigma))
assert sigma > 0, sigma
joints[i, :, 3] = sigma
return joints
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 test_bottomup_pipeline():
data_prefix = 'tests/data/coco/'
ann_file = osp.join(data_prefix, 'test_coco.json')
coco = COCO(ann_file)
ann_info = {}
ann_info['flip_pairs'] = [[1, 2], [3, 4], [5, 6], [7, 8], [9, 10],
[11, 12], [13, 14], [15, 16]]
ann_info['flip_index'] = [
0, 2, 1, 4, 3, 6, 5, 8, 7, 10, 9, 12, 11, 14, 13, 16, 15
]
ann_info['use_different_joint_weights'] = False
ann_info['joint_weights'] = np.array([
1., 1., 1., 1., 1., 1., 1., 1.2, 1.2, 1.5, 1.5, 1., 1., 1.2, 1.2, 1.5,
1.5
],
dtype=np.float32).reshape((17, 1))
ann_info['image_size'] = np.array([384, 512])
ann_info['heatmap_size'] = np.array([[96, 128], [192, 256]])
ann_info['num_joints'] = 17
ann_info['num_scales'] = 2
ann_info['scale_aware_sigma'] = False
ann_ids = coco.getAnnIds(785)
anno = coco.loadAnns(ann_ids)
mask = _get_mask(coco, anno, 785)
anno = [
obj for obj in anno if obj['iscrowd'] == 0 or obj['num_keypoints'] > 0
]
joints = _get_joints(anno, ann_info, False)
mask_list = [mask.copy() for _ in range(ann_info['num_scales'])]
joints_list = [joints.copy() for _ in range(ann_info['num_scales'])]
results = {}
results['dataset'] = 'coco'
results['image_file'] = osp.join(data_prefix, '000000000785.jpg')
results['mask'] = mask_list
results['joints'] = joints_list
results['ann_info'] = ann_info
transform = LoadImageFromFile()
results = transform(copy.deepcopy(results))
assert results['img'].shape == (425, 640, 3)
# test HorizontalFlip
random_horizontal_flip = BottomUpRandomFlip(flip_prob=1.)
results_horizontal_flip = random_horizontal_flip(copy.deepcopy(results))
assert _check_flip(results['img'], results_horizontal_flip['img'])
random_horizontal_flip = BottomUpRandomFlip(flip_prob=0.)
results_horizontal_flip = random_horizontal_flip(copy.deepcopy(results))
assert (results['img'] == results_horizontal_flip['img']).all()
results_copy = copy.deepcopy(results)
results_copy['mask'] = mask_list[0]
with pytest.raises(AssertionError):
results_horizontal_flip = random_horizontal_flip(
copy.deepcopy(results_copy))
results_copy = copy.deepcopy(results)
results_copy['joints'] = joints_list[0]
with pytest.raises(AssertionError):
results_horizontal_flip = random_horizontal_flip(
copy.deepcopy(results_copy))
results_copy = copy.deepcopy(results)
results_copy['joints'] = joints_list[:1]
with pytest.raises(AssertionError):
results_horizontal_flip = random_horizontal_flip(
copy.deepcopy(results_copy))
results_copy = copy.deepcopy(results)
results_copy['mask'] = mask_list[:1]
with pytest.raises(AssertionError):
results_horizontal_flip = random_horizontal_flip(
copy.deepcopy(results_copy))
# test TopDownAffine
random_affine_transform = BottomUpRandomAffine(30, [0.75, 1.5], 'short', 0)
results_affine_transform = random_affine_transform(copy.deepcopy(results))
assert results_affine_transform['img'].shape == (512, 384, 3)
random_affine_transform = BottomUpRandomAffine(30, [0.75, 1.5], 'short',
40)
results_affine_transform = random_affine_transform(copy.deepcopy(results))
assert results_affine_transform['img'].shape == (512, 384, 3)
results_copy = copy.deepcopy(results)
results_copy['ann_info']['scale_aware_sigma'] = True
joints = _get_joints(anno, results_copy['ann_info'], False)
results_copy['joints'] = \
[joints.copy() for _ in range(results_copy['ann_info']['num_scales'])]
results_affine_transform = random_affine_transform(results_copy)
assert results_affine_transform['img'].shape == (512, 384, 3)
results_copy = copy.deepcopy(results)
results_copy['mask'] = mask_list[0]
with pytest.raises(AssertionError):
results_horizontal_flip = random_affine_transform(
copy.deepcopy(results_copy))
results_copy = copy.deepcopy(results)
results_copy['joints'] = joints_list[0]
with pytest.raises(AssertionError):
results_horizontal_flip = random_affine_transform(
copy.deepcopy(results_copy))
results_copy = copy.deepcopy(results)
results_copy['joints'] = joints_list[:1]
with pytest.raises(AssertionError):
results_horizontal_flip = random_affine_transform(
copy.deepcopy(results_copy))
results_copy = copy.deepcopy(results)
results_copy['mask'] = mask_list[:1]
with pytest.raises(AssertionError):
results_horizontal_flip = random_affine_transform(
copy.deepcopy(results_copy))
random_affine_transform = BottomUpRandomAffine(30, [0.75, 1.5], 'long', 40)
results_affine_transform = random_affine_transform(copy.deepcopy(results))
assert results_affine_transform['img'].shape == (512, 384, 3)
with pytest.raises(ValueError):
random_affine_transform = BottomUpRandomAffine(30, [0.75, 1.5],
'short-long', 40)
results_affine_transform = random_affine_transform(
copy.deepcopy(results))
# test BottomUpGenerateTarget
generate_multi_target = BottomUpGenerateTarget(2, 30)
results_generate_multi_target = generate_multi_target(
copy.deepcopy(results))
assert 'targets' in results_generate_multi_target
assert len(results_generate_multi_target['targets']
) == results['ann_info']['num_scales']
# test BottomUpGetImgSize when W > H
get_multi_scale_size = BottomUpGetImgSize([1])
results_get_multi_scale_size = get_multi_scale_size(copy.deepcopy(results))
assert 'test_scale_factor' in results_get_multi_scale_size['ann_info']
assert 'base_size' in results_get_multi_scale_size['ann_info']
assert 'center' in results_get_multi_scale_size['ann_info']
assert 'scale' in results_get_multi_scale_size['ann_info']
assert results_get_multi_scale_size['ann_info']['base_size'][1] == 512
# test BottomUpResizeAlign
transforms = [
dict(type='ToTensor'),
dict(
type='NormalizeTensor',
mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225]),
]
resize_align_multi_scale = BottomUpResizeAlign(transforms=transforms)
results_copy = copy.deepcopy(results_get_multi_scale_size)
results_resize_align_multi_scale = resize_align_multi_scale(results_copy)
assert 'aug_data' in results_resize_align_multi_scale['ann_info']
# test when W < H
ann_info['image_size'] = np.array([512, 384])
ann_info['heatmap_size'] = np.array([[128, 96], [256, 192]])
results = {}
results['dataset'] = 'coco'
results['image_file'] = osp.join(data_prefix, '000000000785.jpg')
results['mask'] = mask_list
results['joints'] = joints_list
results['ann_info'] = ann_info
results['img'] = np.random.rand(640, 425, 3)
# test HorizontalFlip
random_horizontal_flip = BottomUpRandomFlip(flip_prob=1.)
results_horizontal_flip = random_horizontal_flip(copy.deepcopy(results))
assert _check_flip(results['img'], results_horizontal_flip['img'])
random_horizontal_flip = BottomUpRandomFlip(flip_prob=0.)
results_horizontal_flip = random_horizontal_flip(copy.deepcopy(results))
assert (results['img'] == results_horizontal_flip['img']).all()
results_copy = copy.deepcopy(results)
results_copy['mask'] = mask_list[0]
with pytest.raises(AssertionError):
results_horizontal_flip = random_horizontal_flip(
copy.deepcopy(results_copy))
results_copy = copy.deepcopy(results)
results_copy['joints'] = joints_list[0]
with pytest.raises(AssertionError):
results_horizontal_flip = random_horizontal_flip(
copy.deepcopy(results_copy))
results_copy = copy.deepcopy(results)
results_copy['joints'] = joints_list[:1]
with pytest.raises(AssertionError):
results_horizontal_flip = random_horizontal_flip(
copy.deepcopy(results_copy))
results_copy = copy.deepcopy(results)
results_copy['mask'] = mask_list[:1]
with pytest.raises(AssertionError):
results_horizontal_flip = random_horizontal_flip(
copy.deepcopy(results_copy))
# test TopDownAffine
random_affine_transform = BottomUpRandomAffine(30, [0.75, 1.5], 'short', 0)
results_affine_transform = random_affine_transform(copy.deepcopy(results))
assert results_affine_transform['img'].shape == (384, 512, 3)
random_affine_transform = BottomUpRandomAffine(30, [0.75, 1.5], 'short',
40)
results_affine_transform = random_affine_transform(copy.deepcopy(results))
assert results_affine_transform['img'].shape == (384, 512, 3)
results_copy = copy.deepcopy(results)
results_copy['ann_info']['scale_aware_sigma'] = True
joints = _get_joints(anno, results_copy['ann_info'], False)
results_copy['joints'] = \
[joints.copy() for _ in range(results_copy['ann_info']['num_scales'])]
results_affine_transform = random_affine_transform(results_copy)
assert results_affine_transform['img'].shape == (384, 512, 3)
results_copy = copy.deepcopy(results)
results_copy['mask'] = mask_list[0]
with pytest.raises(AssertionError):
results_horizontal_flip = random_affine_transform(
copy.deepcopy(results_copy))
results_copy = copy.deepcopy(results)
results_copy['joints'] = joints_list[0]
with pytest.raises(AssertionError):
results_horizontal_flip = random_affine_transform(
copy.deepcopy(results_copy))
results_copy = copy.deepcopy(results)
results_copy['joints'] = joints_list[:1]
with pytest.raises(AssertionError):
results_horizontal_flip = random_affine_transform(
copy.deepcopy(results_copy))
results_copy = copy.deepcopy(results)
results_copy['mask'] = mask_list[:1]
with pytest.raises(AssertionError):
results_horizontal_flip = random_affine_transform(
copy.deepcopy(results_copy))
random_affine_transform = BottomUpRandomAffine(30, [0.75, 1.5], 'long', 40)
results_affine_transform = random_affine_transform(copy.deepcopy(results))
assert results_affine_transform['img'].shape == (384, 512, 3)
with pytest.raises(ValueError):
random_affine_transform = BottomUpRandomAffine(30, [0.75, 1.5],
'short-long', 40)
results_affine_transform = random_affine_transform(
copy.deepcopy(results))
# test BottomUpGenerateTarget
generate_multi_target = BottomUpGenerateTarget(2, 30)
results_generate_multi_target = generate_multi_target(
copy.deepcopy(results))
assert 'targets' in results_generate_multi_target
assert len(results_generate_multi_target['targets']
) == results['ann_info']['num_scales']
# test BottomUpGetImgSize when W < H
get_multi_scale_size = BottomUpGetImgSize([1])
results_get_multi_scale_size = get_multi_scale_size(copy.deepcopy(results))
assert 'test_scale_factor' in results_get_multi_scale_size['ann_info']
assert 'base_size' in results_get_multi_scale_size['ann_info']
assert 'center' in results_get_multi_scale_size['ann_info']
assert 'scale' in results_get_multi_scale_size['ann_info']
assert results_get_multi_scale_size['ann_info']['base_size'][0] == 512
def test_BottomUpGenerateHeatmapTarget():
data_prefix = 'tests/data/coco/'
ann_file = osp.join(data_prefix, 'test_coco.json')
coco = COCO(ann_file)
ann_info = {}
ann_info['heatmap_size'] = np.array([128, 256])
ann_info['num_joints'] = 17
ann_info['num_scales'] = 2
ann_info['scale_aware_sigma'] = False
ann_ids = coco.getAnnIds(785)
anno = coco.loadAnns(ann_ids)
mask = _get_mask(coco, anno, 785)
anno = [
obj for obj in anno if obj['iscrowd'] == 0 or obj['num_keypoints'] > 0
]
joints = _get_joints(anno, ann_info, False)
mask_list = [mask.copy() for _ in range(ann_info['num_scales'])]
joints_list = [joints.copy() for _ in range(ann_info['num_scales'])]
results = {}
results['dataset'] = 'coco'
results['image_file'] = osp.join(data_prefix, '000000000785.jpg')
results['mask'] = mask_list
results['joints'] = joints_list
results['ann_info'] = ann_info
generate_heatmap_target = BottomUpGenerateHeatmapTarget(2)
results_generate_heatmap_target = generate_heatmap_target(results)
assert 'target' in results_generate_heatmap_target
assert len(results_generate_heatmap_target['target']
) == results['ann_info']['num_scales']
def test_BottomUpGeneratePAFTarget():
ann_info = {}
ann_info['skeleton'] = [[0, 1], [2, 3]]
ann_info['heatmap_size'] = np.array([5])
ann_info['num_joints'] = 4
ann_info['num_scales'] = 1
mask = np.ones((5, 5), dtype=bool)
joints = np.array([[[1, 1, 2], [3, 3, 2], [0, 0, 0], [0, 0, 0]],
[[1, 3, 2], [3, 1, 2], [0, 0, 0], [0, 0, 0]]])
mask_list = [mask.copy() for _ in range(ann_info['num_scales'])]
joints_list = [joints.copy() for _ in range(ann_info['num_scales'])]
results = {}
results['dataset'] = 'coco'
results['mask'] = mask_list
results['joints'] = joints_list
results['ann_info'] = ann_info
generate_paf_target = BottomUpGeneratePAFTarget(1)
results_generate_paf_target = generate_paf_target(results)
sqrt = np.sqrt(2) / 2
assert (results_generate_paf_target['target'] == np.array(
[[[sqrt, sqrt, 0, sqrt, sqrt], [sqrt, sqrt, sqrt, sqrt, sqrt],
[0, sqrt, sqrt, sqrt, 0], [sqrt, sqrt, sqrt, sqrt, sqrt],
[sqrt, sqrt, 0, sqrt, sqrt]],
[[sqrt, sqrt, 0, -sqrt, -sqrt], [sqrt, sqrt, 0, -sqrt, -sqrt],
[0, 0, 0, 0, 0], [-sqrt, -sqrt, 0, sqrt, sqrt],
[-sqrt, -sqrt, 0, sqrt, sqrt]],
[[0, 0, 0, 0, 0], [0, 0, 0, 0, 0], [0, 0, 0, 0, 0], [0, 0, 0, 0, 0],
[0, 0, 0, 0, 0]],
[[0, 0, 0, 0, 0], [0, 0, 0, 0, 0], [0, 0, 0, 0, 0], [0, 0, 0, 0, 0],
[0, 0, 0, 0, 0]]],
dtype=np.float32)).all()