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.
129 lines
4.0 KiB
129 lines
4.0 KiB
# Copyright (c) OpenMMLab. All rights reserved.
|
|
import tempfile
|
|
|
|
from mmcv import Config
|
|
|
|
from mmpose.datasets import DATASETS, build_dataloader
|
|
from mmpose.models import builder
|
|
|
|
|
|
def test_voxelpose_forward():
|
|
dataset = 'Body3DMviewDirectPanopticDataset'
|
|
dataset_class = DATASETS.get(dataset)
|
|
dataset_info = Config.fromfile(
|
|
'configs/_base_/datasets/panoptic_body3d.py').dataset_info
|
|
space_size = [8000, 8000, 2000]
|
|
space_center = [0, -500, 800]
|
|
cube_size = [20, 20, 8]
|
|
data_cfg = dict(
|
|
image_size=[960, 512],
|
|
heatmap_size=[[240, 128]],
|
|
space_size=space_size,
|
|
space_center=space_center,
|
|
cube_size=cube_size,
|
|
num_joints=15,
|
|
seq_list=['160906_band1'],
|
|
cam_list=[(0, 12), (0, 6)],
|
|
num_cameras=2,
|
|
seq_frame_interval=1,
|
|
subset='train',
|
|
need_2d_label=True,
|
|
need_camera_param=True,
|
|
root_id=2)
|
|
|
|
pipeline = [
|
|
dict(
|
|
type='MultiItemProcess',
|
|
pipeline=[
|
|
dict(
|
|
type='BottomUpGenerateTarget', sigma=3, max_num_people=20)
|
|
]),
|
|
dict(
|
|
type='DiscardDuplicatedItems',
|
|
keys_list=[
|
|
'joints_3d', 'joints_3d_visible', 'ann_info', 'roots_3d',
|
|
'num_persons', 'sample_id'
|
|
]),
|
|
dict(
|
|
type='GenerateVoxel3DHeatmapTarget',
|
|
sigma=200.0,
|
|
joint_indices=[2]),
|
|
dict(type='RenameKeys', key_pairs=[('targets', 'input_heatmaps')]),
|
|
dict(
|
|
type='Collect',
|
|
keys=['targets_3d', 'input_heatmaps'],
|
|
meta_keys=[
|
|
'camera', 'center', 'scale', 'joints_3d', 'num_persons',
|
|
'joints_3d_visible', 'roots_3d', 'sample_id'
|
|
]),
|
|
]
|
|
|
|
model_cfg = dict(
|
|
type='DetectAndRegress',
|
|
backbone=None,
|
|
human_detector=dict(
|
|
type='VoxelCenterDetector',
|
|
image_size=[960, 512],
|
|
heatmap_size=[240, 128],
|
|
space_size=space_size,
|
|
cube_size=cube_size,
|
|
space_center=space_center,
|
|
center_net=dict(
|
|
type='V2VNet', input_channels=15, output_channels=1),
|
|
center_head=dict(
|
|
type='CuboidCenterHead',
|
|
space_size=space_size,
|
|
space_center=space_center,
|
|
cube_size=cube_size,
|
|
max_num=3,
|
|
max_pool_kernel=3),
|
|
train_cfg=dict(dist_threshold=500000000.0),
|
|
test_cfg=dict(center_threshold=0.0),
|
|
),
|
|
pose_regressor=dict(
|
|
type='VoxelSinglePose',
|
|
image_size=[960, 512],
|
|
heatmap_size=[240, 128],
|
|
sub_space_size=[2000, 2000, 2000],
|
|
sub_cube_size=[20, 20, 8],
|
|
num_joints=15,
|
|
pose_net=dict(
|
|
type='V2VNet', input_channels=15, output_channels=15),
|
|
pose_head=dict(type='CuboidPoseHead', beta=100.0),
|
|
train_cfg=None,
|
|
test_cfg=None))
|
|
|
|
model = builder.build_posenet(model_cfg)
|
|
with tempfile.TemporaryDirectory() as tmpdir:
|
|
dataset = dataset_class(
|
|
ann_file=tmpdir + '/tmp_train.pkl',
|
|
img_prefix='tests/data/panoptic_body3d/',
|
|
data_cfg=data_cfg,
|
|
pipeline=pipeline,
|
|
dataset_info=dataset_info,
|
|
test_mode=False)
|
|
|
|
data_loader = build_dataloader(
|
|
dataset,
|
|
seed=None,
|
|
dist=False,
|
|
shuffle=False,
|
|
drop_last=False,
|
|
workers_per_gpu=1,
|
|
samples_per_gpu=1)
|
|
|
|
for data in data_loader:
|
|
# test forward_train
|
|
_ = model(
|
|
img=None,
|
|
img_metas=data['img_metas'].data[0],
|
|
return_loss=True,
|
|
targets_3d=data['targets_3d'],
|
|
input_heatmaps=data['input_heatmaps'])
|
|
|
|
# test forward_test
|
|
_ = model(
|
|
img=None,
|
|
img_metas=data['img_metas'].data[0],
|
|
return_loss=False,
|
|
input_heatmaps=data['input_heatmaps'])
|
|
|