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.
258 lines
9.6 KiB
258 lines
9.6 KiB
# Copyright (c) OpenMMLab. All rights reserved.
|
|
import os
|
|
import os.path as osp
|
|
from argparse import ArgumentParser
|
|
|
|
import mmcv
|
|
import numpy as np
|
|
from xtcocotools.coco import COCO
|
|
|
|
from mmpose.apis import inference_interhand_3d_model, vis_3d_pose_result
|
|
from mmpose.apis.inference import init_pose_model
|
|
from mmpose.core import SimpleCamera
|
|
|
|
|
|
def _transform_interhand_camera_param(interhand_camera_param):
|
|
"""Transform the camera parameters in interhand2.6m dataset to the format
|
|
of SimpleCamera.
|
|
|
|
Args:
|
|
interhand_camera_param (dict): camera parameters including:
|
|
- camrot: 3x3, camera rotation matrix (world-to-camera)
|
|
- campos: 3x1, camera location in world space
|
|
- focal: 2x1, camera focal length
|
|
- princpt: 2x1, camera center
|
|
|
|
Returns:
|
|
param (dict): camera parameters including:
|
|
- R: 3x3, camera rotation matrix (camera-to-world)
|
|
- T: 3x1, camera translation (camera-to-world)
|
|
- f: 2x1, camera focal length
|
|
- c: 2x1, camera center
|
|
"""
|
|
camera_param = {}
|
|
camera_param['R'] = np.array(interhand_camera_param['camrot']).T
|
|
camera_param['T'] = np.array(interhand_camera_param['campos'])[:, None]
|
|
camera_param['f'] = np.array(interhand_camera_param['focal'])[:, None]
|
|
camera_param['c'] = np.array(interhand_camera_param['princpt'])[:, None]
|
|
return camera_param
|
|
|
|
|
|
def main():
|
|
parser = ArgumentParser()
|
|
parser.add_argument('pose_config', help='Config file for pose network')
|
|
parser.add_argument('pose_checkpoint', help='Checkpoint file')
|
|
parser.add_argument('--img-root', type=str, default='', help='Image root')
|
|
parser.add_argument(
|
|
'--json-file',
|
|
type=str,
|
|
default='',
|
|
help='Json file containing image info.')
|
|
parser.add_argument(
|
|
'--camera-param-file',
|
|
type=str,
|
|
default=None,
|
|
help='Camera parameter file for converting 3D pose predictions from '
|
|
' the pixel space to camera space. If None, keypoints in pixel space'
|
|
'will be visualized')
|
|
parser.add_argument(
|
|
'--gt-joints-file',
|
|
type=str,
|
|
default=None,
|
|
help='Optional argument. Ground truth 3D keypoint parameter file. '
|
|
'If None, gt keypoints will not be shown and keypoints in pixel '
|
|
'space will be visualized.')
|
|
parser.add_argument(
|
|
'--rebase-keypoint-height',
|
|
action='store_true',
|
|
help='Rebase the predicted 3D pose so its lowest keypoint has a '
|
|
'height of 0 (landing on the ground). This is useful for '
|
|
'visualization when the model do not predict the global position '
|
|
'of the 3D pose.')
|
|
parser.add_argument(
|
|
'--show-ground-truth',
|
|
action='store_true',
|
|
help='If True, show ground truth keypoint if it is available.')
|
|
parser.add_argument(
|
|
'--show',
|
|
action='store_true',
|
|
default=False,
|
|
help='whether to show img')
|
|
parser.add_argument(
|
|
'--out-img-root',
|
|
type=str,
|
|
default=None,
|
|
help='Root of the output visualization images. '
|
|
'Default not saving the visualization images.')
|
|
parser.add_argument(
|
|
'--device', default='cuda:0', help='Device for inference')
|
|
parser.add_argument(
|
|
'--kpt-thr', type=float, default=0.3, help='Keypoint score threshold')
|
|
parser.add_argument(
|
|
'--radius',
|
|
type=int,
|
|
default=4,
|
|
help='Keypoint radius for visualization')
|
|
parser.add_argument(
|
|
'--thickness',
|
|
type=int,
|
|
default=1,
|
|
help='Link thickness for visualization')
|
|
|
|
args = parser.parse_args()
|
|
assert args.show or (args.out_img_root != '')
|
|
|
|
coco = COCO(args.json_file)
|
|
|
|
# build the pose model from a config file and a checkpoint file
|
|
pose_model = init_pose_model(
|
|
args.pose_config, args.pose_checkpoint, device=args.device.lower())
|
|
dataset = pose_model.cfg.data['test']['type']
|
|
|
|
# load camera parameters
|
|
camera_params = None
|
|
if args.camera_param_file is not None:
|
|
camera_params = mmcv.load(args.camera_param_file)
|
|
# load ground truth joints parameters
|
|
gt_joint_params = None
|
|
if args.gt_joints_file is not None:
|
|
gt_joint_params = mmcv.load(args.gt_joints_file)
|
|
|
|
# load hand bounding boxes
|
|
det_results_list = []
|
|
for image_id, image in coco.imgs.items():
|
|
image_name = osp.join(args.img_root, image['file_name'])
|
|
|
|
ann_ids = coco.getAnnIds(image_id)
|
|
det_results = []
|
|
|
|
capture_key = str(image['capture'])
|
|
camera_key = image['camera']
|
|
frame_idx = image['frame_idx']
|
|
|
|
for ann_id in ann_ids:
|
|
ann = coco.anns[ann_id]
|
|
if camera_params is not None:
|
|
camera_param = {
|
|
key: camera_params[capture_key][key][camera_key]
|
|
for key in camera_params[capture_key].keys()
|
|
}
|
|
camera_param = _transform_interhand_camera_param(camera_param)
|
|
else:
|
|
camera_param = None
|
|
if gt_joint_params is not None:
|
|
joint_param = gt_joint_params[capture_key][str(frame_idx)]
|
|
gt_joint = np.concatenate([
|
|
np.array(joint_param['world_coord']),
|
|
np.array(joint_param['joint_valid'])
|
|
],
|
|
axis=-1)
|
|
else:
|
|
gt_joint = None
|
|
|
|
det_result = {
|
|
'image_name': image_name,
|
|
'bbox': ann['bbox'], # bbox format is 'xywh'
|
|
'camera_param': camera_param,
|
|
'keypoints_3d_gt': gt_joint
|
|
}
|
|
det_results.append(det_result)
|
|
det_results_list.append(det_results)
|
|
|
|
for i, det_results in enumerate(
|
|
mmcv.track_iter_progress(det_results_list)):
|
|
|
|
image_name = det_results[0]['image_name']
|
|
|
|
pose_results = inference_interhand_3d_model(
|
|
pose_model, image_name, det_results, dataset=dataset)
|
|
|
|
# Post processing
|
|
pose_results_vis = []
|
|
for idx, res in enumerate(pose_results):
|
|
keypoints_3d = res['keypoints_3d']
|
|
# normalize kpt score
|
|
if keypoints_3d[:, 3].max() > 1:
|
|
keypoints_3d[:, 3] /= 255
|
|
# get 2D keypoints in pixel space
|
|
res['keypoints'] = keypoints_3d[:, [0, 1, 3]]
|
|
|
|
# For model-predicted keypoints, channel 0 and 1 are coordinates
|
|
# in pixel space, and channel 2 is the depth (in mm) relative
|
|
# to root joints.
|
|
# If both camera parameter and absolute depth of root joints are
|
|
# provided, we can transform keypoint to camera space for better
|
|
# visualization.
|
|
camera_param = res['camera_param']
|
|
keypoints_3d_gt = res['keypoints_3d_gt']
|
|
if camera_param is not None and keypoints_3d_gt is not None:
|
|
# build camera model
|
|
camera = SimpleCamera(camera_param)
|
|
# transform gt joints from world space to camera space
|
|
keypoints_3d_gt[:, :3] = camera.world_to_camera(
|
|
keypoints_3d_gt[:, :3])
|
|
|
|
# transform relative depth to absolute depth
|
|
keypoints_3d[:21, 2] += keypoints_3d_gt[20, 2]
|
|
keypoints_3d[21:, 2] += keypoints_3d_gt[41, 2]
|
|
|
|
# transform keypoints from pixel space to camera space
|
|
keypoints_3d[:, :3] = camera.pixel_to_camera(
|
|
keypoints_3d[:, :3])
|
|
|
|
# rotate the keypoint to make z-axis correspondent to height
|
|
# for better visualization
|
|
vis_R = np.array([[1, 0, 0], [0, 0, -1], [0, 1, 0]])
|
|
keypoints_3d[:, :3] = keypoints_3d[:, :3] @ vis_R
|
|
if keypoints_3d_gt is not None:
|
|
keypoints_3d_gt[:, :3] = keypoints_3d_gt[:, :3] @ vis_R
|
|
|
|
# rebase height (z-axis)
|
|
if args.rebase_keypoint_height:
|
|
valid = keypoints_3d[..., 3] > 0
|
|
keypoints_3d[..., 2] -= np.min(
|
|
keypoints_3d[valid, 2], axis=-1, keepdims=True)
|
|
res['keypoints_3d'] = keypoints_3d
|
|
res['keypoints_3d_gt'] = keypoints_3d_gt
|
|
|
|
# Add title
|
|
instance_id = res.get('track_id', idx)
|
|
res['title'] = f'Prediction ({instance_id})'
|
|
pose_results_vis.append(res)
|
|
# Add ground truth
|
|
if args.show_ground_truth:
|
|
if keypoints_3d_gt is None:
|
|
print('Fail to show ground truth. Please make sure that'
|
|
' gt-joints-file is provided.')
|
|
else:
|
|
gt = res.copy()
|
|
if args.rebase_keypoint_height:
|
|
valid = keypoints_3d_gt[..., 3] > 0
|
|
keypoints_3d_gt[..., 2] -= np.min(
|
|
keypoints_3d_gt[valid, 2], axis=-1, keepdims=True)
|
|
gt['keypoints_3d'] = keypoints_3d_gt
|
|
gt['title'] = f'Ground truth ({instance_id})'
|
|
pose_results_vis.append(gt)
|
|
|
|
# Visualization
|
|
if args.out_img_root is None:
|
|
out_file = None
|
|
else:
|
|
os.makedirs(args.out_img_root, exist_ok=True)
|
|
out_file = osp.join(args.out_img_root, f'vis_{i}.jpg')
|
|
|
|
vis_3d_pose_result(
|
|
pose_model,
|
|
result=pose_results_vis,
|
|
img=det_results[0]['image_name'],
|
|
out_file=out_file,
|
|
dataset=dataset,
|
|
show=args.show,
|
|
kpt_score_thr=args.kpt_thr,
|
|
radius=args.radius,
|
|
thickness=args.thickness,
|
|
)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main()
|
|
|