mirror of https://github.com/open-mmlab/mmpose
554 lines
21 KiB
Python
554 lines
21 KiB
Python
# Copyright (c) OpenMMLab. All rights reserved.
|
|
import logging
|
|
import mimetypes
|
|
import os
|
|
import time
|
|
from argparse import ArgumentParser
|
|
from functools import partial
|
|
|
|
import cv2
|
|
import json_tricks as json
|
|
import mmcv
|
|
import mmengine
|
|
import numpy as np
|
|
from mmengine.logging import print_log
|
|
|
|
from mmpose.apis import (_track_by_iou, _track_by_oks,
|
|
convert_keypoint_definition, extract_pose_sequence,
|
|
inference_pose_lifter_model, inference_topdown,
|
|
init_model)
|
|
from mmpose.models.pose_estimators import PoseLifter
|
|
from mmpose.models.pose_estimators.topdown import TopdownPoseEstimator
|
|
from mmpose.registry import VISUALIZERS
|
|
from mmpose.structures import (PoseDataSample, merge_data_samples,
|
|
split_instances)
|
|
from mmpose.utils import adapt_mmdet_pipeline
|
|
|
|
try:
|
|
from mmdet.apis import inference_detector, init_detector
|
|
has_mmdet = True
|
|
except (ImportError, ModuleNotFoundError):
|
|
has_mmdet = False
|
|
|
|
|
|
def parse_args():
|
|
parser = ArgumentParser()
|
|
parser.add_argument('det_config', help='Config file for detection')
|
|
parser.add_argument('det_checkpoint', help='Checkpoint file for detection')
|
|
parser.add_argument(
|
|
'pose_estimator_config',
|
|
type=str,
|
|
default=None,
|
|
help='Config file for the 1st stage 2D pose estimator')
|
|
parser.add_argument(
|
|
'pose_estimator_checkpoint',
|
|
type=str,
|
|
default=None,
|
|
help='Checkpoint file for the 1st stage 2D pose estimator')
|
|
parser.add_argument(
|
|
'pose_lifter_config',
|
|
help='Config file for the 2nd stage pose lifter model')
|
|
parser.add_argument(
|
|
'pose_lifter_checkpoint',
|
|
help='Checkpoint file for the 2nd stage pose lifter model')
|
|
parser.add_argument('--input', type=str, default='', help='Video path')
|
|
parser.add_argument(
|
|
'--show',
|
|
action='store_true',
|
|
default=False,
|
|
help='Whether to show visualizations')
|
|
parser.add_argument(
|
|
'--disable-rebase-keypoint',
|
|
action='store_true',
|
|
default=False,
|
|
help='Whether to disable rebasing the predicted 3D pose so its '
|
|
'lowest keypoint has a height of 0 (landing on the ground). Rebase '
|
|
'is useful for visualization when the model do not predict the '
|
|
'global position of the 3D pose.')
|
|
parser.add_argument(
|
|
'--disable-norm-pose-2d',
|
|
action='store_true',
|
|
default=False,
|
|
help='Whether to scale the bbox (along with the 2D pose) to the '
|
|
'average bbox scale of the dataset, and move the bbox (along with the '
|
|
'2D pose) to the average bbox center of the dataset. This is useful '
|
|
'when bbox is small, especially in multi-person scenarios.')
|
|
parser.add_argument(
|
|
'--num-instances',
|
|
type=int,
|
|
default=1,
|
|
help='The number of 3D poses to be visualized in every frame. If '
|
|
'less than 0, it will be set to the number of pose results in the '
|
|
'first frame.')
|
|
parser.add_argument(
|
|
'--output-root',
|
|
type=str,
|
|
default='',
|
|
help='Root of the output video file. '
|
|
'Default not saving the visualization video.')
|
|
parser.add_argument(
|
|
'--save-predictions',
|
|
action='store_true',
|
|
default=False,
|
|
help='Whether to save predicted results')
|
|
parser.add_argument(
|
|
'--device', default='cuda:0', help='Device used for inference')
|
|
parser.add_argument(
|
|
'--det-cat-id',
|
|
type=int,
|
|
default=0,
|
|
help='Category id for bounding box detection model')
|
|
parser.add_argument(
|
|
'--bbox-thr',
|
|
type=float,
|
|
default=0.3,
|
|
help='Bounding box score threshold')
|
|
parser.add_argument('--kpt-thr', type=float, default=0.3)
|
|
parser.add_argument(
|
|
'--use-oks-tracking', action='store_true', help='Using OKS tracking')
|
|
parser.add_argument(
|
|
'--tracking-thr', type=float, default=0.3, help='Tracking threshold')
|
|
parser.add_argument(
|
|
'--show-interval', type=int, default=0, help='Sleep seconds per frame')
|
|
parser.add_argument(
|
|
'--thickness',
|
|
type=int,
|
|
default=1,
|
|
help='Link thickness for visualization')
|
|
parser.add_argument(
|
|
'--radius',
|
|
type=int,
|
|
default=3,
|
|
help='Keypoint radius for visualization')
|
|
parser.add_argument(
|
|
'--online',
|
|
action='store_true',
|
|
default=False,
|
|
help='Inference mode. If set to True, can not use future frame'
|
|
'information when using multi frames for inference in the 2D pose'
|
|
'detection stage. Default: False.')
|
|
|
|
args = parser.parse_args()
|
|
return args
|
|
|
|
|
|
def process_one_image(args, detector, frame, frame_idx, pose_estimator,
|
|
pose_est_results_last, pose_est_results_list, next_id,
|
|
pose_lifter, visualize_frame, visualizer):
|
|
"""Visualize detected and predicted keypoints of one image.
|
|
|
|
Pipeline of this function:
|
|
|
|
frame
|
|
|
|
|
V
|
|
+-----------------+
|
|
| detector |
|
|
+-----------------+
|
|
| det_result
|
|
V
|
|
+-----------------+
|
|
| pose_estimator |
|
|
+-----------------+
|
|
| pose_est_results
|
|
V
|
|
+--------------------------------------------+
|
|
| convert 2d kpts into pose-lifting format |
|
|
+--------------------------------------------+
|
|
| pose_est_results_list
|
|
V
|
|
+-----------------------+
|
|
| extract_pose_sequence |
|
|
+-----------------------+
|
|
| pose_seq_2d
|
|
V
|
|
+-------------+
|
|
| pose_lifter |
|
|
+-------------+
|
|
| pose_lift_results
|
|
V
|
|
+-----------------+
|
|
| post-processing |
|
|
+-----------------+
|
|
| pred_3d_data_samples
|
|
V
|
|
+------------+
|
|
| visualizer |
|
|
+------------+
|
|
|
|
Args:
|
|
args (Argument): Custom command-line arguments.
|
|
detector (mmdet.BaseDetector): The mmdet detector.
|
|
frame (np.ndarray): The image frame read from input image or video.
|
|
frame_idx (int): The index of current frame.
|
|
pose_estimator (TopdownPoseEstimator): The pose estimator for 2d pose.
|
|
pose_est_results_last (list(PoseDataSample)): The results of pose
|
|
estimation from the last frame for tracking instances.
|
|
pose_est_results_list (list(list(PoseDataSample))): The list of all
|
|
pose estimation results converted by
|
|
``convert_keypoint_definition`` from previous frames. In
|
|
pose-lifting stage it is used to obtain the 2d estimation sequence.
|
|
next_id (int): The next track id to be used.
|
|
pose_lifter (PoseLifter): The pose-lifter for estimating 3d pose.
|
|
visualize_frame (np.ndarray): The image for drawing the results on.
|
|
visualizer (Visualizer): The visualizer for visualizing the 2d and 3d
|
|
pose estimation results.
|
|
|
|
Returns:
|
|
pose_est_results (list(PoseDataSample)): The pose estimation result of
|
|
the current frame.
|
|
pose_est_results_list (list(list(PoseDataSample))): The list of all
|
|
converted pose estimation results until the current frame.
|
|
pred_3d_instances (InstanceData): The result of pose-lifting.
|
|
Specifically, the predicted keypoints and scores are saved at
|
|
``pred_3d_instances.keypoints`` and
|
|
``pred_3d_instances.keypoint_scores``.
|
|
next_id (int): The next track id to be used.
|
|
"""
|
|
pose_lift_dataset = pose_lifter.cfg.test_dataloader.dataset
|
|
pose_lift_dataset_name = pose_lifter.dataset_meta['dataset_name']
|
|
|
|
# First stage: conduct 2D pose detection in a Topdown manner
|
|
# use detector to obtain person bounding boxes
|
|
det_result = inference_detector(detector, frame)
|
|
pred_instance = det_result.pred_instances.cpu().numpy()
|
|
|
|
# filter out the person instances with category and bbox threshold
|
|
# e.g. 0 for person in COCO
|
|
bboxes = pred_instance.bboxes
|
|
bboxes = bboxes[np.logical_and(pred_instance.labels == args.det_cat_id,
|
|
pred_instance.scores > args.bbox_thr)]
|
|
|
|
# estimate pose results for current image
|
|
pose_est_results = inference_topdown(pose_estimator, frame, bboxes)
|
|
|
|
if args.use_oks_tracking:
|
|
_track = partial(_track_by_oks)
|
|
else:
|
|
_track = _track_by_iou
|
|
|
|
pose_det_dataset_name = pose_estimator.dataset_meta['dataset_name']
|
|
pose_est_results_converted = []
|
|
|
|
# convert 2d pose estimation results into the format for pose-lifting
|
|
# such as changing the keypoint order, flipping the keypoint, etc.
|
|
for i, data_sample in enumerate(pose_est_results):
|
|
pred_instances = data_sample.pred_instances.cpu().numpy()
|
|
keypoints = pred_instances.keypoints
|
|
# calculate area and bbox
|
|
if 'bboxes' in pred_instances:
|
|
areas = np.array([(bbox[2] - bbox[0]) * (bbox[3] - bbox[1])
|
|
for bbox in pred_instances.bboxes])
|
|
pose_est_results[i].pred_instances.set_field(areas, 'areas')
|
|
else:
|
|
areas, bboxes = [], []
|
|
for keypoint in keypoints:
|
|
xmin = np.min(keypoint[:, 0][keypoint[:, 0] > 0], initial=1e10)
|
|
xmax = np.max(keypoint[:, 0])
|
|
ymin = np.min(keypoint[:, 1][keypoint[:, 1] > 0], initial=1e10)
|
|
ymax = np.max(keypoint[:, 1])
|
|
areas.append((xmax - xmin) * (ymax - ymin))
|
|
bboxes.append([xmin, ymin, xmax, ymax])
|
|
pose_est_results[i].pred_instances.areas = np.array(areas)
|
|
pose_est_results[i].pred_instances.bboxes = np.array(bboxes)
|
|
|
|
# track id
|
|
track_id, pose_est_results_last, _ = _track(data_sample,
|
|
pose_est_results_last,
|
|
args.tracking_thr)
|
|
if track_id == -1:
|
|
if np.count_nonzero(keypoints[:, :, 1]) >= 3:
|
|
track_id = next_id
|
|
next_id += 1
|
|
else:
|
|
# If the number of keypoints detected is small,
|
|
# delete that person instance.
|
|
keypoints[:, :, 1] = -10
|
|
pose_est_results[i].pred_instances.set_field(
|
|
keypoints, 'keypoints')
|
|
pose_est_results[i].pred_instances.set_field(
|
|
pred_instances.bboxes * 0, 'bboxes')
|
|
pose_est_results[i].set_field(pred_instances, 'pred_instances')
|
|
track_id = -1
|
|
pose_est_results[i].set_field(track_id, 'track_id')
|
|
|
|
# convert keypoints for pose-lifting
|
|
pose_est_result_converted = PoseDataSample()
|
|
pose_est_result_converted.set_field(
|
|
pose_est_results[i].pred_instances.clone(), 'pred_instances')
|
|
pose_est_result_converted.set_field(
|
|
pose_est_results[i].gt_instances.clone(), 'gt_instances')
|
|
keypoints = convert_keypoint_definition(keypoints,
|
|
pose_det_dataset_name,
|
|
pose_lift_dataset_name)
|
|
pose_est_result_converted.pred_instances.set_field(
|
|
keypoints, 'keypoints')
|
|
pose_est_result_converted.set_field(pose_est_results[i].track_id,
|
|
'track_id')
|
|
pose_est_results_converted.append(pose_est_result_converted)
|
|
|
|
pose_est_results_list.append(pose_est_results_converted.copy())
|
|
|
|
# Second stage: Pose lifting
|
|
# extract and pad input pose2d sequence
|
|
pose_seq_2d = extract_pose_sequence(
|
|
pose_est_results_list,
|
|
frame_idx=frame_idx,
|
|
causal=pose_lift_dataset.get('causal', False),
|
|
seq_len=pose_lift_dataset.get('seq_len', 1),
|
|
step=pose_lift_dataset.get('seq_step', 1))
|
|
|
|
# conduct 2D-to-3D pose lifting
|
|
norm_pose_2d = not args.disable_norm_pose_2d
|
|
pose_lift_results = inference_pose_lifter_model(
|
|
pose_lifter,
|
|
pose_seq_2d,
|
|
image_size=visualize_frame.shape[:2],
|
|
norm_pose_2d=norm_pose_2d)
|
|
|
|
# post-processing
|
|
for idx, pose_lift_result in enumerate(pose_lift_results):
|
|
pose_lift_result.track_id = pose_est_results[idx].get('track_id', 1e4)
|
|
|
|
pred_instances = pose_lift_result.pred_instances
|
|
keypoints = pred_instances.keypoints
|
|
keypoint_scores = pred_instances.keypoint_scores
|
|
if keypoint_scores.ndim == 3:
|
|
keypoint_scores = np.squeeze(keypoint_scores, axis=1)
|
|
pose_lift_results[
|
|
idx].pred_instances.keypoint_scores = keypoint_scores
|
|
if keypoints.ndim == 4:
|
|
keypoints = np.squeeze(keypoints, axis=1)
|
|
|
|
keypoints = keypoints[..., [0, 2, 1]]
|
|
keypoints[..., 0] = -keypoints[..., 0]
|
|
keypoints[..., 2] = -keypoints[..., 2]
|
|
|
|
# rebase height (z-axis)
|
|
if not args.disable_rebase_keypoint:
|
|
keypoints[..., 2] -= np.min(
|
|
keypoints[..., 2], axis=-1, keepdims=True)
|
|
|
|
pose_lift_results[idx].pred_instances.keypoints = keypoints
|
|
|
|
pose_lift_results = sorted(
|
|
pose_lift_results, key=lambda x: x.get('track_id', 1e4))
|
|
|
|
pred_3d_data_samples = merge_data_samples(pose_lift_results)
|
|
det_data_sample = merge_data_samples(pose_est_results)
|
|
pred_3d_instances = pred_3d_data_samples.get('pred_instances', None)
|
|
|
|
if args.num_instances < 0:
|
|
args.num_instances = len(pose_lift_results)
|
|
|
|
# Visualization
|
|
if visualizer is not None:
|
|
visualizer.add_datasample(
|
|
'result',
|
|
visualize_frame,
|
|
data_sample=pred_3d_data_samples,
|
|
det_data_sample=det_data_sample,
|
|
draw_gt=False,
|
|
dataset_2d=pose_det_dataset_name,
|
|
dataset_3d=pose_lift_dataset_name,
|
|
show=args.show,
|
|
draw_bbox=True,
|
|
kpt_thr=args.kpt_thr,
|
|
num_instances=args.num_instances,
|
|
wait_time=args.show_interval)
|
|
|
|
return pose_est_results, pose_est_results_list, pred_3d_instances, next_id
|
|
|
|
|
|
def main():
|
|
assert has_mmdet, 'Please install mmdet to run the demo.'
|
|
|
|
args = parse_args()
|
|
|
|
assert args.show or (args.output_root != '')
|
|
assert args.input != ''
|
|
assert args.det_config is not None
|
|
assert args.det_checkpoint is not None
|
|
|
|
detector = init_detector(
|
|
args.det_config, args.det_checkpoint, device=args.device.lower())
|
|
detector.cfg = adapt_mmdet_pipeline(detector.cfg)
|
|
|
|
pose_estimator = init_model(
|
|
args.pose_estimator_config,
|
|
args.pose_estimator_checkpoint,
|
|
device=args.device.lower())
|
|
|
|
assert isinstance(pose_estimator, TopdownPoseEstimator), 'Only "TopDown"' \
|
|
'model is supported for the 1st stage (2D pose detection)'
|
|
|
|
det_kpt_color = pose_estimator.dataset_meta.get('keypoint_colors', None)
|
|
det_dataset_skeleton = pose_estimator.dataset_meta.get(
|
|
'skeleton_links', None)
|
|
det_dataset_link_color = pose_estimator.dataset_meta.get(
|
|
'skeleton_link_colors', None)
|
|
|
|
pose_lifter = init_model(
|
|
args.pose_lifter_config,
|
|
args.pose_lifter_checkpoint,
|
|
device=args.device.lower())
|
|
|
|
assert isinstance(pose_lifter, PoseLifter), \
|
|
'Only "PoseLifter" model is supported for the 2nd stage ' \
|
|
'(2D-to-3D lifting)'
|
|
|
|
pose_lifter.cfg.visualizer.radius = args.radius
|
|
pose_lifter.cfg.visualizer.line_width = args.thickness
|
|
pose_lifter.cfg.visualizer.det_kpt_color = det_kpt_color
|
|
pose_lifter.cfg.visualizer.det_dataset_skeleton = det_dataset_skeleton
|
|
pose_lifter.cfg.visualizer.det_dataset_link_color = det_dataset_link_color
|
|
visualizer = VISUALIZERS.build(pose_lifter.cfg.visualizer)
|
|
|
|
# the dataset_meta is loaded from the checkpoint
|
|
visualizer.set_dataset_meta(pose_lifter.dataset_meta)
|
|
|
|
if args.input == 'webcam':
|
|
input_type = 'webcam'
|
|
else:
|
|
input_type = mimetypes.guess_type(args.input)[0].split('/')[0]
|
|
|
|
if args.output_root == '':
|
|
save_output = False
|
|
else:
|
|
mmengine.mkdir_or_exist(args.output_root)
|
|
output_file = os.path.join(args.output_root,
|
|
os.path.basename(args.input))
|
|
if args.input == 'webcam':
|
|
output_file += '.mp4'
|
|
save_output = True
|
|
|
|
if args.save_predictions:
|
|
assert args.output_root != ''
|
|
args.pred_save_path = f'{args.output_root}/results_' \
|
|
f'{os.path.splitext(os.path.basename(args.input))[0]}.json'
|
|
|
|
if save_output:
|
|
fourcc = cv2.VideoWriter_fourcc(*'mp4v')
|
|
|
|
pose_est_results_list = []
|
|
pred_instances_list = []
|
|
if input_type == 'image':
|
|
frame = mmcv.imread(args.input, channel_order='rgb')
|
|
_, _, pred_3d_instances, _ = process_one_image(
|
|
args=args,
|
|
detector=detector,
|
|
frame=frame,
|
|
frame_idx=0,
|
|
pose_estimator=pose_estimator,
|
|
pose_est_results_last=[],
|
|
pose_est_results_list=pose_est_results_list,
|
|
next_id=0,
|
|
pose_lifter=pose_lifter,
|
|
visualize_frame=frame,
|
|
visualizer=visualizer)
|
|
|
|
if args.save_predictions:
|
|
# save prediction results
|
|
pred_instances_list = split_instances(pred_3d_instances)
|
|
|
|
if save_output:
|
|
frame_vis = visualizer.get_image()
|
|
mmcv.imwrite(mmcv.rgb2bgr(frame_vis), output_file)
|
|
|
|
elif input_type in ['webcam', 'video']:
|
|
next_id = 0
|
|
pose_est_results = []
|
|
|
|
if args.input == 'webcam':
|
|
video = cv2.VideoCapture(0)
|
|
else:
|
|
video = cv2.VideoCapture(args.input)
|
|
|
|
(major_ver, minor_ver, subminor_ver) = (cv2.__version__).split('.')
|
|
if int(major_ver) < 3:
|
|
fps = video.get(cv2.cv.CV_CAP_PROP_FPS)
|
|
else:
|
|
fps = video.get(cv2.CAP_PROP_FPS)
|
|
|
|
video_writer = None
|
|
frame_idx = 0
|
|
|
|
while video.isOpened():
|
|
success, frame = video.read()
|
|
frame_idx += 1
|
|
|
|
if not success:
|
|
break
|
|
|
|
pose_est_results_last = pose_est_results
|
|
|
|
# First stage: 2D pose detection
|
|
# make person results for current image
|
|
(pose_est_results, pose_est_results_list, pred_3d_instances,
|
|
next_id) = process_one_image(
|
|
args=args,
|
|
detector=detector,
|
|
frame=frame,
|
|
frame_idx=frame_idx,
|
|
pose_estimator=pose_estimator,
|
|
pose_est_results_last=pose_est_results_last,
|
|
pose_est_results_list=pose_est_results_list,
|
|
next_id=next_id,
|
|
pose_lifter=pose_lifter,
|
|
visualize_frame=mmcv.bgr2rgb(frame),
|
|
visualizer=visualizer)
|
|
|
|
if args.save_predictions:
|
|
# save prediction results
|
|
pred_instances_list.append(
|
|
dict(
|
|
frame_id=frame_idx,
|
|
instances=split_instances(pred_3d_instances)))
|
|
|
|
if save_output:
|
|
frame_vis = visualizer.get_image()
|
|
if video_writer is None:
|
|
# the size of the image with visualization may vary
|
|
# depending on the presence of heatmaps
|
|
video_writer = cv2.VideoWriter(output_file, fourcc, fps,
|
|
(frame_vis.shape[1],
|
|
frame_vis.shape[0]))
|
|
|
|
video_writer.write(mmcv.rgb2bgr(frame_vis))
|
|
|
|
if args.show:
|
|
# press ESC to exit
|
|
if cv2.waitKey(5) & 0xFF == 27:
|
|
break
|
|
time.sleep(args.show_interval)
|
|
|
|
video.release()
|
|
|
|
if video_writer:
|
|
video_writer.release()
|
|
else:
|
|
args.save_predictions = False
|
|
raise ValueError(
|
|
f'file {os.path.basename(args.input)} has invalid format.')
|
|
|
|
if args.save_predictions:
|
|
with open(args.pred_save_path, 'w') as f:
|
|
json.dump(
|
|
dict(
|
|
meta_info=pose_lifter.dataset_meta,
|
|
instance_info=pred_instances_list),
|
|
f,
|
|
indent='\t')
|
|
print(f'predictions have been saved at {args.pred_save_path}')
|
|
|
|
if save_output:
|
|
input_type = input_type.replace('webcam', 'video')
|
|
print_log(
|
|
f'the output {input_type} has been saved at {output_file}',
|
|
logger='current',
|
|
level=logging.INFO)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main()
|