mmpose-0.29.0/tests/test_apis/test_inference_3d.py

# Copyright (c) OpenMMLab. All rights reserved.
import os.path as osp
import tempfile

import mmcv
import numpy as np
import pytest
import torch

from mmpose.apis import (extract_pose_sequence, inference_interhand_3d_model,
 inference_mesh_model, inference_pose_lifter_model,
 init_pose_model, vis_3d_mesh_result,
 vis_3d_pose_result)
from mmpose.datasets.dataset_info import DatasetInfo
from tests.utils.mesh_utils import generate_smpl_weight_file


def test_pose_lifter_demo():
 # H36M demo
 pose_model = init_pose_model(
 'configs/body/3d_kpt_sview_rgb_img/pose_lift/'
 'h36m/simplebaseline3d_h36m.py',
 None,
 device='cpu')

 pose_det_result = {
 'keypoints': np.zeros((17, 3)),
 'bbox': [50, 50, 50, 50],
 'track_id': 0,
 'image_name': 'tests/data/h36m/S1_Directions_1.54138969_000001.jpg',
 }

 pose_results_2d = [[pose_det_result]]

 dataset_info = DatasetInfo(pose_model.cfg.data['test']['dataset_info'])

 pose_results_2d = extract_pose_sequence(
 pose_results_2d, frame_idx=0, causal=False, seq_len=1, step=1)

 _ = inference_pose_lifter_model(
 pose_model,
 pose_results_2d,
 dataset_info=dataset_info,
 with_track_id=False)

 pose_lift_results = inference_pose_lifter_model(
 pose_model,
 pose_results_2d,
 dataset_info=dataset_info,
 with_track_id=True)

 for res in pose_lift_results:
 res['title'] = 'title'
 vis_3d_pose_result(
 pose_model,
 pose_lift_results,
 img=pose_results_2d[0][0]['image_name'],
 dataset_info=dataset_info)

 # test special cases
 # Empty 2D results
 _ = inference_pose_lifter_model(
 pose_model, [[]], dataset_info=dataset_info, with_track_id=False)

 if torch.cuda.is_available():
 _ = inference_pose_lifter_model(
 pose_model.cuda(),
 pose_results_2d,
 dataset_info=dataset_info,
 with_track_id=False)

 # test videopose3d
 pose_model = init_pose_model(
 'configs/body/3d_kpt_sview_rgb_vid/video_pose_lift/h36m/'
 'videopose3d_h36m_243frames_fullconv_supervised_cpn_ft.py',
 None,
 device='cpu')

 pose_det_result_0 = {
 'keypoints': np.ones((17, 3)),
 'bbox': [50, 50, 100, 100],
 'track_id': 0,
 'image_name': 'tests/data/h36m/S1_Directions_1.54138969_000001.jpg',
 }
 pose_det_result_1 = {
 'keypoints': np.ones((17, 3)),
 'bbox': [50, 50, 100, 100],
 'track_id': 1,
 'image_name': 'tests/data/h36m/S5_SittingDown.54138969_002061.jpg',
 }
 pose_det_result_2 = {
 'keypoints': np.ones((17, 3)),
 'bbox': [50, 50, 100, 100],
 'track_id': 2,
 'image_name': 'tests/data/h36m/S7_Greeting.55011271_000396.jpg',
 }

 pose_results_2d = [[pose_det_result_0], [pose_det_result_1],
 [pose_det_result_2]]

 dataset_info = DatasetInfo(pose_model.cfg.data['test']['dataset_info'])

 seq_len = pose_model.cfg.test_data_cfg.seq_len
 pose_results_2d_seq = extract_pose_sequence(
 pose_results_2d, 1, causal=False, seq_len=seq_len, step=1)

 pose_lift_results = inference_pose_lifter_model(
 pose_model,
 pose_results_2d_seq,
 dataset_info=dataset_info,
 with_track_id=True,
 image_size=[1000, 1000],
 norm_pose_2d=True)

 for res in pose_lift_results:
 res['title'] = 'title'
 vis_3d_pose_result(
 pose_model,
 pose_lift_results,
 img=pose_results_2d[0][0]['image_name'],
 dataset_info=dataset_info,
 )


def test_interhand3d_demo():
 # H36M demo
 pose_model = init_pose_model(
 'configs/hand/3d_kpt_sview_rgb_img/internet/interhand3d/'
 'res50_interhand3d_all_256x256.py',
 None,
 device='cpu')

 image_name = 'tests/data/interhand2.6m/image2017.jpg'
 det_result = {
 'image_name': image_name,
 'bbox': [50, 50, 50, 50], # bbox format is 'xywh'
 'camera_param': None,
 'keypoints_3d_gt': None
 }
 det_results = [det_result]
 dataset = pose_model.cfg.data['test']['type']
 dataset_info = DatasetInfo(pose_model.cfg.data['test']['dataset_info'])

 pose_results = inference_interhand_3d_model(
 pose_model, image_name, det_results, dataset=dataset)

 for res in pose_results:
 res['title'] = 'title'

 vis_3d_pose_result(
 pose_model,
 result=pose_results,
 img=det_results[0]['image_name'],
 dataset_info=dataset_info,
 )

 # test special cases
 # Empty det results
 _ = inference_interhand_3d_model(
 pose_model, image_name, [], dataset=dataset)

 if torch.cuda.is_available():
 _ = inference_interhand_3d_model(
 pose_model.cuda(), image_name, det_results, dataset=dataset)

 with pytest.raises(NotImplementedError):
 _ = inference_interhand_3d_model(
 pose_model, image_name, det_results, dataset='test')


def test_body_mesh_demo():
 # H36M demo
 config = 'configs/body/3d_mesh_sview_rgb_img/hmr' \
 '/mixed/res50_mixed_224x224.py'
 config = mmcv.Config.fromfile(config)
 config.model.mesh_head.smpl_mean_params = \
 'tests/data/smpl/smpl_mean_params.npz'

 pose_model = None
 with tempfile.TemporaryDirectory() as tmpdir:
 config.model.smpl.smpl_path = tmpdir
 config.model.smpl.joints_regressor = osp.join(
 tmpdir, 'test_joint_regressor.npy')
 # generate weight file for SMPL model.
 generate_smpl_weight_file(tmpdir)
 pose_model = init_pose_model(config, device='cpu')

 assert pose_model is not None, 'Fail to build pose model'

 image_name = 'tests/data/h36m/S1_Directions_1.54138969_000001.jpg'
 det_result = {
 'keypoints': np.zeros((17, 3)),
 'bbox': [50, 50, 50, 50],
 'image_name': image_name,
 }

 # make person bounding boxes
 person_results = [det_result]
 dataset = pose_model.cfg.data['test']['type']

 # test a single image, with a list of bboxes
 pose_results = inference_mesh_model(
 pose_model,
 image_name,
 person_results,
 bbox_thr=None,
 format='xywh',
 dataset=dataset)

 vis_3d_mesh_result(pose_model, pose_results, image_name)