def test_bottom_up_demo():
# build the pose model from a config file and a checkpoint file
pose_model = init_pose_model(
'configs/body/2d_kpt_sview_rgb_img/associative_embedding/'
'coco/res50_coco_512x512.py',
None,
device='cpu')
image_name = 'tests/data/coco/000000000785.jpg'
dataset_info = DatasetInfo(pose_model.cfg.data['test'].get(
'dataset_info', None))
pose_results, _ = inference_bottom_up_pose_model(pose_model,
image_name,
dataset_info=dataset_info)
# show the results
vis_pose_result(pose_model,
image_name,
pose_results,
dataset_info=dataset_info)
# test dataset_info without sigmas
pose_model_copy = copy.deepcopy(pose_model)
pose_model_copy.cfg.data.test.dataset_info.pop('sigmas')
pose_results, _ = inference_bottom_up_pose_model(pose_model_copy,
image_name,
dataset_info=dataset_info)
def __init__(self,
ann_file,
img_prefix,
data_cfg,
pipeline,
dataset_info=None,
test_mode=False):
self.image_info = {}
self.ann_info = {}
self.ann_file = ann_file
self.img_prefix = img_prefix
self.pipeline = pipeline
self.test_mode = test_mode
self.ann_info['image_size'] = np.array(data_cfg['image_size'])
self.ann_info['heatmap_size'] = np.array(data_cfg['heatmap_size'])
self.ann_info['num_joints'] = data_cfg['num_joints']
self.ann_info['space_size'] = data_cfg['space_size']
self.ann_info['space_center'] = data_cfg['space_center']
self.ann_info['cube_size'] = data_cfg['cube_size']
self.ann_info['scale_aware_sigma'] = data_cfg.get(
'scale_aware_sigma', False)
if dataset_info is None:
raise ValueError(
'Check https://github.com/open-mmlab/mmpose/pull/663 '
'for details.')
dataset_info = DatasetInfo(dataset_info)
assert self.ann_info['num_joints'] <= dataset_info.keypoint_num
self.ann_info['flip_pairs'] = dataset_info.flip_pairs
self.ann_info['num_scales'] = 1
self.ann_info['flip_index'] = dataset_info.flip_index
self.ann_info['upper_body_ids'] = dataset_info.upper_body_ids
self.ann_info['lower_body_ids'] = dataset_info.lower_body_ids
self.ann_info['joint_weights'] = dataset_info.joint_weights
self.ann_info['skeleton'] = dataset_info.skeleton
self.sigmas = dataset_info.sigmas
self.dataset_name = dataset_info.dataset_name
self.load_config(data_cfg)
self.db = []
self.pipeline = Compose(self.pipeline)
def __init__(self,
ann_file,
img_prefix,
data_cfg,
pipeline,
dataset_info=None,
test_mode=False):
self.ann_file = ann_file
self.img_prefix = img_prefix
self.data_cfg = copy.deepcopy(data_cfg)
self.pipeline = pipeline
self.test_mode = test_mode
self.ann_info = {}
if dataset_info is None:
raise ValueError(
'Check https://github.com/open-mmlab/mmpose/pull/663 '
'for details.')
dataset_info = DatasetInfo(dataset_info)
self.load_config(self.data_cfg)
self.ann_info['num_joints'] = data_cfg['num_joints']
assert self.ann_info['num_joints'] == dataset_info.keypoint_num
self.ann_info['flip_pairs'] = dataset_info.flip_pairs
self.ann_info['upper_body_ids'] = dataset_info.upper_body_ids
self.ann_info['lower_body_ids'] = dataset_info.lower_body_ids
self.ann_info['joint_weights'] = dataset_info.joint_weights
self.ann_info['skeleton'] = dataset_info.skeleton
self.sigmas = dataset_info.sigmas
self.dataset_name = dataset_info.dataset_name
self.data_info = self.load_annotations()
self.sample_indices = self.build_sample_indices()
self.pipeline = Compose(pipeline)
self.name2id = {
name: i
for i, name in enumerate(self.data_info['imgnames'])
}
def _draw_keypoint(self, canvas: np.ndarray, input_msg: FrameMessage):
"""Draw object keypoints."""
objects = input_msg.get_objects(lambda x: 'pose_model_cfg' in x)
# return if there is no object with keypoints
if not objects:
return canvas
for model_cfg, group in groupby(objects,
lambda x: x['pose_model_cfg']):
dataset_info = DatasetInfo(model_cfg.dataset_info)
keypoints = [obj['keypoints'] for obj in group]
imshow_keypoints(canvas,
keypoints,
skeleton=dataset_info.skeleton,
kpt_score_thr=0.3,
pose_kpt_color=dataset_info.pose_kpt_color,
pose_link_color=dataset_info.pose_link_color,
radius=self.radius,
thickness=self.thickness)
return canvas
def main():
"""Visualize the demo images.
Using mmdet to detect the human.
"""
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_config', help='Config file for pose')
parser.add_argument('pose_checkpoint', help='Checkpoint file for pose')
parser.add_argument('--img-root', type=str, default='', help='Image root')
parser.add_argument('--img', type=str, default='', help='Image file')
parser.add_argument(
'--show',
action='store_true',
default=False,
help='whether to show img')
parser.add_argument(
'--out-img-root',
type=str,
default='',
help='root of the output img file. '
'Default not saving the visualization images.')
parser.add_argument(
'--device', default='cuda:0', help='Device used for inference')
parser.add_argument(
'--det-cat-id',
type=int,
default=1,
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, 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')
assert has_mmdet, 'Please install mmdet to run the demo.'
args = parser.parse_args()
assert args.show or (args.out_img_root != '')
assert args.img != ''
assert args.det_config is not None
assert args.det_checkpoint is not None
det_model = init_detector(
args.det_config, args.det_checkpoint, device=args.device.lower())
# 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']
dataset_info = pose_model.cfg.data['test'].get('dataset_info', None)
if dataset_info is None:
warnings.warn(
'Please set `dataset_info` in the config.'
'Check https://github.com/open-mmlab/mmpose/pull/663 for details.',
DeprecationWarning)
else:
dataset_info = DatasetInfo(dataset_info)
image_name = os.path.join(args.img_root, args.img)
# test a single image, the resulting box is (x1, y1, x2, y2)
mmdet_results = inference_detector(det_model, image_name)
# keep the person class bounding boxes.
person_results = process_mmdet_results(mmdet_results, args.det_cat_id)
# test a single image, with a list of bboxes.
# optional
return_heatmap = False
# e.g. use ('backbone', ) to return backbone feature
output_layer_names = None
pose_results, returned_outputs = inference_top_down_pose_model(
pose_model,
image_name,
person_results,
bbox_thr=args.bbox_thr,
format='xyxy',
dataset=dataset,
dataset_info=dataset_info,
return_heatmap=return_heatmap,
outputs=output_layer_names)
if args.out_img_root == '':
out_file = None
else:
os.makedirs(args.out_img_root, exist_ok=True)
out_file = os.path.join(args.out_img_root, f'vis_{args.img}')
# show the results
vis_pose_result(
pose_model,
image_name,
pose_results,
dataset=dataset,
dataset_info=dataset_info,
kpt_score_thr=args.kpt_thr,
radius=args.radius,
thickness=args.thickness,
show=args.show,
out_file=out_file)
def main():
"""Visualize the demo images.
Using mmdet to detect the human.
"""
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_config', help='Config file for pose')
parser.add_argument('pose_checkpoint', help='Checkpoint file for pose')
parser.add_argument('--video-path', type=str, help='Video path')
parser.add_argument('--show',
action='store_true',
default=False,
help='whether to show visualizations.')
parser.add_argument('--out-video-root',
default='',
help='Root of the output video file. '
'Default not saving the visualization video.')
parser.add_argument('--device',
default='cuda:0',
help='Device used for inference')
parser.add_argument('--det-cat-id',
type=int,
default=1,
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,
help='Keypoint score threshold')
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('--euro',
action='store_true',
help='Using One_Euro_Filter for smoothing')
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')
assert has_mmdet, 'Please install mmdet to run the demo.'
args = parser.parse_args()
assert args.show or (args.out_video_root != '')
assert args.det_config is not None
assert args.det_checkpoint is not None
det_model = init_detector(args.det_config,
args.det_checkpoint,
device=args.device.lower())
# 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']
dataset_info = pose_model.cfg.data['test'].get('dataset_info', None)
if dataset_info is None:
warnings.warn(
'Please set `dataset_info` in the config.'
'Check https://github.com/open-mmlab/mmpose/pull/663 for details.',
DeprecationWarning)
else:
dataset_info = DatasetInfo(dataset_info)
cap = cv2.VideoCapture(args.video_path)
fps = None
assert cap.isOpened(), f'Faild to load video file {args.video_path}'
if args.out_video_root == '':
save_out_video = False
else:
os.makedirs(args.out_video_root, exist_ok=True)
save_out_video = True
if save_out_video:
fps = cap.get(cv2.CAP_PROP_FPS)
size = (int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)),
int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)))
fourcc = cv2.VideoWriter_fourcc(*'mp4v')
videoWriter = cv2.VideoWriter(
os.path.join(args.out_video_root,
f'vis_{os.path.basename(args.video_path)}'), fourcc,
fps, size)
# optional
return_heatmap = False
# e.g. use ('backbone', ) to return backbone feature
output_layer_names = None
next_id = 0
pose_results = []
while (cap.isOpened()):
pose_results_last = pose_results
flag, img = cap.read()
if not flag:
break
# test a single image, the resulting box is (x1, y1, x2, y2)
mmdet_results = inference_detector(det_model, img)
# keep the person class bounding boxes.
person_results = process_mmdet_results(mmdet_results, args.det_cat_id)
# test a single image, with a list of bboxes.
pose_results, returned_outputs = inference_top_down_pose_model(
pose_model,
img,
person_results,
bbox_thr=args.bbox_thr,
format='xyxy',
dataset=dataset,
dataset_info=dataset_info,
return_heatmap=return_heatmap,
outputs=output_layer_names)
# get track id for each person instance
pose_results, next_id = get_track_id(pose_results,
pose_results_last,
next_id,
use_oks=args.use_oks_tracking,
tracking_thr=args.tracking_thr,
use_one_euro=args.euro,
fps=fps)
# show the results
vis_img = vis_pose_tracking_result(pose_model,
img,
pose_results,
radius=args.radius,
thickness=args.thickness,
dataset=dataset,
dataset_info=dataset_info,
kpt_score_thr=args.kpt_thr,
show=False)
if args.show:
cv2.imshow('Image', vis_img)
if save_out_video:
videoWriter.write(vis_img)
if args.show and cv2.waitKey(1) & 0xFF == ord('q'):
break
cap.release()
if save_out_video:
videoWriter.release()
if args.show:
cv2.destroyAllWindows()
def main():
"""Visualize the demo images."""
parser = ArgumentParser()
parser.add_argument('pose_config', help='Config file for pose')
parser.add_argument('pose_checkpoint', help='Checkpoint file for pose')
parser.add_argument('--video-path', type=str, help='Video path')
parser.add_argument('--show',
action='store_true',
default=False,
help='whether to show visualizations.')
parser.add_argument('--out-video-root',
default='',
help='Root of the output video file. '
'Default not saving the visualization video.')
parser.add_argument('--device',
default='cuda:0',
help='Device used for inference')
parser.add_argument('--kpt-thr',
type=float,
default=0.5,
help='Keypoint score threshold')
parser.add_argument('--pose-nms-thr',
type=float,
default=0.9,
help='OKS threshold for pose NMS')
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(
'--euro',
action='store_true',
help='(Deprecated, please use --smooth and --smooth-filter-cfg) '
'Using One_Euro_Filter for smoothing.')
parser.add_argument(
'--smooth',
action='store_true',
help='Apply a temporal filter to smooth the pose estimation results. '
'See also --smooth-filter-cfg.')
parser.add_argument(
'--smooth-filter-cfg',
type=str,
default='configs/_base_/filters/one_euro.py',
help='Config file of the filter to smooth the pose estimation '
'results. See also --smooth.')
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_video_root != '')
# 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']
dataset_info = pose_model.cfg.data['test'].get('dataset_info', None)
if dataset_info is None:
warnings.warn(
'Please set `dataset_info` in the config.'
'Check https://github.com/open-mmlab/mmpose/pull/663 for details.',
DeprecationWarning)
assert (dataset == 'BottomUpCocoDataset')
else:
dataset_info = DatasetInfo(dataset_info)
video = mmcv.VideoReader(args.video_path)
assert video.opened, f'Faild to load video file {args.video_path}'
if args.out_video_root == '':
save_out_video = False
else:
os.makedirs(args.out_video_root, exist_ok=True)
save_out_video = True
if save_out_video:
fps = video.fps
size = (video.width, video.height)
fourcc = cv2.VideoWriter_fourcc(*'mp4v')
videoWriter = cv2.VideoWriter(
os.path.join(args.out_video_root,
f'vis_{os.path.basename(args.video_path)}'), fourcc,
fps, size)
# optional
return_heatmap = False
# build pose smoother for temporal refinement
if args.euro:
warnings.warn(
'Argument --euro will be deprecated in the future. '
'Please use --smooth to enable temporal smoothing, and '
'--smooth-filter-cfg to set the filter config.',
DeprecationWarning)
smoother = Smoother(filter_cfg='configs/_base_/filters/one_euro.py',
keypoint_dim=2)
elif args.smooth:
smoother = Smoother(filter_cfg=args.smooth_filter_cfg, keypoint_dim=2)
else:
smoother = None
# e.g. use ('backbone', ) to return backbone feature
output_layer_names = None
next_id = 0
pose_results = []
for cur_frame in mmcv.track_iter_progress(video):
pose_results_last = pose_results
pose_results, _ = inference_bottom_up_pose_model(
pose_model,
cur_frame,
dataset=dataset,
dataset_info=dataset_info,
pose_nms_thr=args.pose_nms_thr,
return_heatmap=return_heatmap,
outputs=output_layer_names)
# get track id for each person instance
pose_results, next_id = get_track_id(pose_results,
pose_results_last,
next_id,
use_oks=args.use_oks_tracking,
tracking_thr=args.tracking_thr,
use_one_euro=args.euro,
fps=fps)
# post-process the pose results with smoother
if smoother:
pose_results = smoother.smooth(pose_results)
# show the results
vis_frame = vis_pose_tracking_result(pose_model,
cur_frame,
pose_results,
radius=args.radius,
thickness=args.thickness,
dataset=dataset,
dataset_info=dataset_info,
kpt_score_thr=args.kpt_thr,
show=False)
if args.show:
cv2.imshow('Image', vis_frame)
if save_out_video:
videoWriter.write(vis_frame)
if args.show and cv2.waitKey(1) & 0xFF == ord('q'):
break
if save_out_video:
videoWriter.release()
if args.show:
cv2.destroyAllWindows()
def main():
"""Visualize the demo images."""
parser = ArgumentParser()
parser.add_argument('pose_config', help='Config file for detection')
parser.add_argument('pose_checkpoint', help='Checkpoint file')
parser.add_argument(
'--img-path',
type=str,
help='Path to an image file or a image folder.')
parser.add_argument(
'--show',
action='store_true',
default=False,
help='whether to show img')
parser.add_argument(
'--out-img-root',
type=str,
default='',
help='Root of the output img file. '
'Default not saving the visualization images.')
parser.add_argument(
'--device', default='cuda:0', help='Device used for inference')
parser.add_argument(
'--kpt-thr', type=float, default=0.3, help='Keypoint score threshold')
parser.add_argument(
'--pose-nms-thr',
type=float,
default=0.9,
help='OKS threshold for pose NMS')
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 != '')
# prepare image list
if osp.isfile(args.img_path):
image_list = [args.img_path]
elif osp.isdir(args.img_path):
image_list = [
osp.join(args.img_path, fn) for fn in os.listdir(args.img_path)
if fn.lower().endswith(('.png', '.jpg', '.jpeg', '.tiff', '.bmp'))
]
else:
raise ValueError('Image path should be an image or image folder.'
f'Got invalid image path: {args.img_path}')
# 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']
dataset_info = pose_model.cfg.data['test'].get('dataset_info', None)
if dataset_info is None:
warnings.warn(
'Please set `dataset_info` in the config.'
'Check https://github.com/open-mmlab/mmpose/pull/663 for details.',
DeprecationWarning)
assert (dataset == 'BottomUpCocoDataset')
else:
dataset_info = DatasetInfo(dataset_info)
# optional
return_heatmap = False
# e.g. use ('backbone', ) to return backbone feature
output_layer_names = None
# process each image
for image_name in mmcv.track_iter_progress(image_list):
# test a single image, with a list of bboxes.
pose_results, returned_outputs = inference_bottom_up_pose_model(
pose_model,
image_name,
dataset=dataset,
dataset_info=dataset_info,
pose_nms_thr=args.pose_nms_thr,
return_heatmap=return_heatmap,
outputs=output_layer_names)
if args.out_img_root == '':
out_file = None
else:
os.makedirs(args.out_img_root, exist_ok=True)
out_file = os.path.join(
args.out_img_root,
f'vis_{osp.splitext(osp.basename(image_name))[0]}.jpg')
# show the results
vis_pose_result(
pose_model,
image_name,
pose_results,
radius=args.radius,
thickness=args.thickness,
dataset=dataset,
dataset_info=dataset_info,
kpt_score_thr=args.kpt_thr,
show=args.show,
out_file=out_file)
def test_top_down_demo():
# COCO demo
# build the pose model from a config file and a checkpoint file
pose_model = init_pose_model(
'configs/body/2d_kpt_sview_rgb_img/topdown_heatmap/'
'coco/res50_coco_256x192.py',
None,
device='cpu')
image_name = 'tests/data/coco/000000000785.jpg'
dataset_info = DatasetInfo(pose_model.cfg.data['test'].get(
'dataset_info', None))
person_result = []
person_result.append({'bbox': [50, 50, 50, 100]})
# test a single image, with a list of bboxes.
pose_results, _ = inference_top_down_pose_model(pose_model,
image_name,
person_result,
format='xywh',
dataset_info=dataset_info)
# show the results
vis_pose_result(pose_model,
image_name,
pose_results,
dataset_info=dataset_info)
# MPII demo
# build the pose model from a config file and a checkpoint file
pose_model = init_pose_model(
'configs/body/2d_kpt_sview_rgb_img/deeppose/'
'mpii/res50_mpii_256x256.py',
None,
device='cpu')
image_name = 'tests/data/mpii/004645041.jpg'
dataset_info = DatasetInfo(pose_model.cfg.data['test'].get(
'dataset_info', None))
person_result = []
person_result.append({'bbox': [50, 50, 50, 100]})
# test a single image, with a list of bboxes.
pose_results, _ = inference_top_down_pose_model(pose_model,
image_name,
person_result,
format='xywh',
dataset_info=dataset_info)
# show the results
vis_pose_result(pose_model,
image_name,
pose_results,
dataset_info=dataset_info)
# AIC demo
pose_model = init_pose_model(
'configs/body/2d_kpt_sview_rgb_img/topdown_heatmap/'
'aic/res50_aic_256x192.py',
None,
device='cpu')
image_name = 'tests/data/aic/054d9ce9201beffc76e5ff2169d2af2f027002ca.jpg'
dataset_info = DatasetInfo(pose_model.cfg.data['test'].get(
'dataset_info', None))
# test a single image, with a list of bboxes.
pose_results, _ = inference_top_down_pose_model(pose_model,
image_name,
person_result,
format='xywh',
dataset_info=dataset_info)
# show the results
vis_pose_result(pose_model,
image_name,
pose_results,
dataset_info=dataset_info)
# OneHand10K demo
# build the pose model from a config file and a checkpoint file
pose_model = init_pose_model(
'configs/hand/2d_kpt_sview_rgb_img/topdown_heatmap/'
'onehand10k/res50_onehand10k_256x256.py',
None,
device='cpu')
image_name = 'tests/data/onehand10k/9.jpg'
dataset_info = DatasetInfo(pose_model.cfg.data['test'].get(
'dataset_info', None))
# test a single image, with a list of bboxes.
pose_results, _ = inference_top_down_pose_model(pose_model,
image_name,
person_result,
format='xywh',
dataset_info=dataset_info)
# show the results
vis_pose_result(pose_model,
image_name,
pose_results,
dataset_info=dataset_info)
# InterHand2DDataset demo
# build the pose model from a config file and a checkpoint file
pose_model = init_pose_model(
'configs/hand/2d_kpt_sview_rgb_img/topdown_heatmap/'
'interhand2d/res50_interhand2d_all_256x256.py',
None,
device='cpu')
image_name = 'tests/data/interhand2.6m/image2017.jpg'
dataset_info = DatasetInfo(pose_model.cfg.data['test'].get(
'dataset_info', None))
# test a single image, with a list of bboxes.
pose_results, _ = inference_top_down_pose_model(pose_model,
image_name,
person_result,
format='xywh',
dataset_info=dataset_info)
# show the results
vis_pose_result(pose_model,
image_name,
pose_results,
dataset_info=dataset_info)
# Face300WDataset demo
# build the pose model from a config file and a checkpoint file
pose_model = init_pose_model(
'configs/face/2d_kpt_sview_rgb_img/topdown_heatmap/'
'300w/res50_300w_256x256.py',
None,
device='cpu')
image_name = 'tests/data/300w/indoor_020.png'
dataset_info = DatasetInfo(pose_model.cfg.data['test'].get(
'dataset_info', None))
# test a single image, with a list of bboxes.
pose_results, _ = inference_top_down_pose_model(pose_model,
image_name,
person_result,
format='xywh',
dataset_info=dataset_info)
# show the results
vis_pose_result(pose_model,
image_name,
pose_results,
dataset_info=dataset_info)
# FaceAFLWDataset demo
# build the pose model from a config file and a checkpoint file
pose_model = init_pose_model(
'configs/face/2d_kpt_sview_rgb_img/topdown_heatmap/'
'aflw/res50_aflw_256x256.py',
None,
device='cpu')
image_name = 'tests/data/aflw/image04476.jpg'
dataset_info = DatasetInfo(pose_model.cfg.data['test'].get(
'dataset_info', None))
# test a single image, with a list of bboxes.
pose_results, _ = inference_top_down_pose_model(pose_model,
image_name,
person_result,
format='xywh',
dataset_info=dataset_info)
# show the results
vis_pose_result(pose_model,
image_name,
pose_results,
dataset_info=dataset_info)
# FaceCOFWDataset demo
# build the pose model from a config file and a checkpoint file
pose_model = init_pose_model(
'configs/face/2d_kpt_sview_rgb_img/topdown_heatmap/'
'cofw/res50_cofw_256x256.py',
None,
device='cpu')
image_name = 'tests/data/cofw/001766.jpg'
dataset_info = DatasetInfo(pose_model.cfg.data['test'].get(
'dataset_info', None))
# test a single image, with a list of bboxes.
pose_results, _ = inference_top_down_pose_model(pose_model,
image_name,
person_result,
format='xywh',
dataset_info=dataset_info)
# show the results
vis_pose_result(pose_model,
image_name,
pose_results,
dataset_info=dataset_info)
# test video pose demo
# PoseWarper + PoseTrack18 demo
# build the pose model from a config file and a checkpoint file
pose_model = init_pose_model(
'configs/body/2d_kpt_sview_rgb_vid/posewarper/posetrack18/'
'hrnet_w48_posetrack18_384x288_posewarper_stage2.py',
None,
device='cpu')
dataset_info = DatasetInfo(pose_model.cfg.data['test'].get(
'dataset_info', None))
person_result = []
# the last value is the confidence of bbox
person_result.append({'bbox': [50, 50, 50, 100, 0.5]})
# test a viedo folder
num_frames = len(pose_model.cfg.data_cfg.frame_weight_test)
video_folder = 'tests/data/posetrack18/videos/000001_mpiinew_test'
frames = sorted(glob(osp.join(video_folder, '*.jpg')))[:num_frames]
cur_frame = frames[0]
# test the frames in the format of image paths
pose_results, _ = inference_top_down_pose_model(pose_model,
frames,
person_result,
format='xywh',
bbox_thr=0.3,
dataset_info=dataset_info)
# show the results
vis_pose_result(pose_model,
cur_frame,
pose_results,
dataset_info=dataset_info)
# test when the person_result is None
pose_results, _ = inference_top_down_pose_model(pose_model,
frames,
person_results=None,
format='xywh',
dataset_info=dataset_info)
# test a video file
video_path = 'tests/data/posetrack18/videos/000001_mpiinew_test/'\
'000001_mpiinew_test.mp4'
video = mmcv.VideoReader(video_path)
# get a sample for test
cur_frame = video[0]
frames = video[:num_frames]
person_result = []
person_result.append({'bbox': [50, 75, 100, 150, 0.6]})
# test the frames in the format of image array
pose_results, _ = inference_top_down_pose_model(pose_model,
frames,
person_result,
bbox_thr=0.9,
format='xyxy',
dataset_info=dataset_info)
# # test the frames in the format of image array
pose_results, _ = inference_top_down_pose_model(pose_model,
frames,
person_results=None,
format='xyxy',
dataset_info=dataset_info)
def main():
"""Visualize the demo images."""
parser = ArgumentParser()
parser.add_argument('pose_config', help='Config file for pose')
parser.add_argument('pose_checkpoint', help='Checkpoint file for pose')
parser.add_argument('--video-path', type=str, help='Video path')
parser.add_argument('--show',
action='store_true',
default=False,
help='whether to show visualizations.')
parser.add_argument('--out-video-root',
default='',
help='Root of the output video file. '
'Default not saving the visualization video.')
parser.add_argument('--device',
default='cuda:0',
help='Device used for inference')
parser.add_argument('--kpt-thr',
type=float,
default=0.3,
help='Keypoint score threshold')
parser.add_argument('--pose-nms-thr',
type=float,
default=0.9,
help='OKS threshold for pose NMS')
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_video_root != '')
# 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']
dataset_info = pose_model.cfg.data['test'].get('dataset_info', None)
if dataset_info is None:
warnings.warn(
'Please set `dataset_info` in the config.'
'Check https://github.com/open-mmlab/mmpose/pull/663 for details.',
DeprecationWarning)
assert (dataset == 'BottomUpCocoDataset')
else:
dataset_info = DatasetInfo(dataset_info)
# read video
video = mmcv.VideoReader(args.video_path)
assert video.opened, f'Faild to load video file {args.video_path}'
if args.out_video_root == '':
save_out_video = False
else:
os.makedirs(args.out_video_root, exist_ok=True)
save_out_video = True
if save_out_video:
fps = video.fps
size = (video.width, video.height)
fourcc = cv2.VideoWriter_fourcc(*'mp4v')
videoWriter = cv2.VideoWriter(
os.path.join(args.out_video_root,
f'vis_{os.path.basename(args.video_path)}'), fourcc,
fps, size)
# optional
return_heatmap = False
# e.g. use ('backbone', ) to return backbone feature
output_layer_names = None
print('Running inference...')
for _, cur_frame in enumerate(mmcv.track_iter_progress(video)):
pose_results, _ = inference_bottom_up_pose_model(
pose_model,
cur_frame,
dataset=dataset,
dataset_info=dataset_info,
pose_nms_thr=args.pose_nms_thr,
return_heatmap=return_heatmap,
outputs=output_layer_names)
# show the results
vis_frame = vis_pose_result(pose_model,
cur_frame,
pose_results,
radius=args.radius,
thickness=args.thickness,
dataset=dataset,
dataset_info=dataset_info,
kpt_score_thr=args.kpt_thr,
show=False)
if args.show:
cv2.imshow('Image', vis_frame)
if save_out_video:
videoWriter.write(vis_frame)
if args.show and cv2.waitKey(1) & 0xFF == ord('q'):
break
if save_out_video:
videoWriter.release()
if args.show:
cv2.destroyAllWindows()
def main():
"""Visualize the demo images.
Require the json_file containing boxes.
"""
parser = ArgumentParser()
parser.add_argument('pose_config', help='Config file for detection')
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('--show',
action='store_true',
default=False,
help='whether to show img')
parser.add_argument('--out-img-root',
type=str,
default='',
help='Root of the output img file. '
'Default not saving the visualization images.')
parser.add_argument('--device',
default='cuda:0',
help='Device used 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']
dataset_info = pose_model.cfg.data['test'].get('dataset_info', None)
if dataset_info is None:
warnings.warn(
'Please set `dataset_info` in the config.'
'Check https://github.com/open-mmlab/mmpose/pull/663 for details.',
DeprecationWarning)
else:
dataset_info = DatasetInfo(dataset_info)
img_keys = list(coco.imgs.keys())
# optional
return_heatmap = False
# e.g. use ('backbone', ) to return backbone feature
output_layer_names = None
# process each image
for i in mmcv.track_iter_progress(range(len(img_keys))):
# get bounding box annotations
image_id = img_keys[i]
image = coco.loadImgs(image_id)[0]
image_name = os.path.join(args.img_root, image['file_name'])
ann_ids = coco.getAnnIds(image_id)
# make person bounding boxes
person_results = []
for ann_id in ann_ids:
person = {}
ann = coco.anns[ann_id]
# bbox format is 'xywh'
person['bbox'] = ann['bbox']
person_results.append(person)
# test a single image, with a list of bboxes
pose_results, returned_outputs = inference_top_down_pose_model(
pose_model,
image_name,
person_results,
bbox_thr=None,
format='xywh',
dataset=dataset,
dataset_info=dataset_info,
return_heatmap=return_heatmap,
outputs=output_layer_names)
if args.out_img_root == '':
out_file = None
else:
os.makedirs(args.out_img_root, exist_ok=True)
out_file = os.path.join(args.out_img_root, f'vis_{i}.jpg')
vis_pose_result(pose_model,
image_name,
pose_results,
dataset=dataset,
dataset_info=dataset_info,
kpt_score_thr=args.kpt_thr,
radius=args.radius,
thickness=args.thickness,
show=args.show,
out_file=out_file)
def test_top_down_demo():
# COCO demo
# build the pose model from a config file and a checkpoint file
pose_model = init_pose_model(
'configs/body/2d_kpt_sview_rgb_img/topdown_heatmap/'
'coco/res50_coco_256x192.py',
None,
device='cpu')
image_name = 'tests/data/coco/000000000785.jpg'
dataset_info = DatasetInfo(pose_model.cfg.data['test'].get(
'dataset_info', None))
person_result = []
person_result.append({'bbox': [50, 50, 50, 100]})
# test a single image, with a list of bboxes.
pose_results, _ = inference_top_down_pose_model(pose_model,
image_name,
person_result,
format='xywh',
dataset_info=dataset_info)
# show the results
vis_pose_result(pose_model,
image_name,
pose_results,
dataset_info=dataset_info)
# AIC demo
pose_model = init_pose_model(
'configs/body/2d_kpt_sview_rgb_img/topdown_heatmap/'
'aic/res50_aic_256x192.py',
None,
device='cpu')
image_name = 'tests/data/aic/054d9ce9201beffc76e5ff2169d2af2f027002ca.jpg'
dataset_info = DatasetInfo(pose_model.cfg.data['test'].get(
'dataset_info', None))
# test a single image, with a list of bboxes.
pose_results, _ = inference_top_down_pose_model(pose_model,
image_name,
person_result,
format='xywh',
dataset_info=dataset_info)
# show the results
vis_pose_result(pose_model,
image_name,
pose_results,
dataset_info=dataset_info)
# OneHand10K demo
# build the pose model from a config file and a checkpoint file
pose_model = init_pose_model(
'configs/hand/2d_kpt_sview_rgb_img/topdown_heatmap/'
'onehand10k/res50_onehand10k_256x256.py',
None,
device='cpu')
image_name = 'tests/data/onehand10k/9.jpg'
dataset_info = DatasetInfo(pose_model.cfg.data['test'].get(
'dataset_info', None))
# test a single image, with a list of bboxes.
pose_results, _ = inference_top_down_pose_model(pose_model,
image_name,
person_result,
format='xywh',
dataset_info=dataset_info)
# show the results
vis_pose_result(pose_model,
image_name,
pose_results,
dataset_info=dataset_info)
# InterHand2DDataset demo
# build the pose model from a config file and a checkpoint file
pose_model = init_pose_model(
'configs/hand/2d_kpt_sview_rgb_img/topdown_heatmap/'
'interhand2d/res50_interhand2d_all_256x256.py',
None,
device='cpu')
image_name = 'tests/data/interhand2.6m/image2017.jpg'
dataset_info = DatasetInfo(pose_model.cfg.data['test'].get(
'dataset_info', None))
# test a single image, with a list of bboxes.
pose_results, _ = inference_top_down_pose_model(pose_model,
image_name,
person_result,
format='xywh',
dataset_info=dataset_info)
# show the results
vis_pose_result(pose_model,
image_name,
pose_results,
dataset_info=dataset_info)
# Face300WDataset demo
# build the pose model from a config file and a checkpoint file
pose_model = init_pose_model(
'configs/face/2d_kpt_sview_rgb_img/topdown_heatmap/'
'300w/res50_300w_256x256.py',
None,
device='cpu')
image_name = 'tests/data/300w/indoor_020.png'
dataset_info = DatasetInfo(pose_model.cfg.data['test'].get(
'dataset_info', None))
# test a single image, with a list of bboxes.
pose_results, _ = inference_top_down_pose_model(pose_model,
image_name,
person_result,
format='xywh',
dataset_info=dataset_info)
# show the results
vis_pose_result(pose_model,
image_name,
pose_results,
dataset_info=dataset_info)
# FaceAFLWDataset demo
# build the pose model from a config file and a checkpoint file
pose_model = init_pose_model(
'configs/face/2d_kpt_sview_rgb_img/topdown_heatmap/'
'aflw/res50_aflw_256x256.py',
None,
device='cpu')
image_name = 'tests/data/aflw/image04476.jpg'
dataset_info = DatasetInfo(pose_model.cfg.data['test'].get(
'dataset_info', None))
# test a single image, with a list of bboxes.
pose_results, _ = inference_top_down_pose_model(pose_model,
image_name,
person_result,
format='xywh',
dataset_info=dataset_info)
# show the results
vis_pose_result(pose_model,
image_name,
pose_results,
dataset_info=dataset_info)
# FaceCOFWDataset demo
# build the pose model from a config file and a checkpoint file
pose_model = init_pose_model(
'configs/face/2d_kpt_sview_rgb_img/topdown_heatmap/'
'cofw/res50_cofw_256x256.py',
None,
device='cpu')
image_name = 'tests/data/cofw/001766.jpg'
dataset_info = DatasetInfo(pose_model.cfg.data['test'].get(
'dataset_info', None))
# test a single image, with a list of bboxes.
pose_results, _ = inference_top_down_pose_model(pose_model,
image_name,
person_result,
format='xywh',
dataset_info=dataset_info)
# show the results
vis_pose_result(pose_model,
image_name,
pose_results,
dataset_info=dataset_info)
def main():
parser = ArgumentParser()
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('--pose-detector-config',
type=str,
default=None,
help='Config file for the 1st stage 2D pose detector')
parser.add_argument(
'--pose-detector-checkpoint',
type=str,
default=None,
help='Checkpoint file for the 1st stage 2D pose detector')
parser.add_argument('--img-root', type=str, default='', help='Image root')
parser.add_argument(
'--json-file',
type=str,
default=None,
help='Json file containing image and bbox information. Optionally,'
'The Json file can also contain 2D pose information. See'
'"only-second-stage"')
parser.add_argument(
'--camera-param-file',
type=str,
default=None,
help='Camera parameter file for converting 3D pose predictions from '
' the camera space to to world space. If None, no conversion will be '
'applied.')
parser.add_argument(
'--only-second-stage',
action='store_true',
help='If true, load 2D pose detection result from the Json file and '
'skip the 1st stage. The pose detection model will be ignored.')
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 if it is available. The ground truth '
'should be contained in the annotations in the Json file with the key '
'"keypoints_3d" for each instance.')
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)
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)
# First stage: 2D pose detection
pose_det_results_list = []
if args.only_second_stage:
from mmpose.apis.inference import _xywh2xyxy
print('Stage 1: load 2D pose results from Json file.')
for image_id, image in coco.imgs.items():
image_name = osp.join(args.img_root, image['file_name'])
ann_ids = coco.getAnnIds(image_id)
pose_det_results = []
for ann_id in ann_ids:
ann = coco.anns[ann_id]
keypoints = np.array(ann['keypoints']).reshape(-1, 3)
keypoints[..., 2] = keypoints[..., 2] >= 1
keypoints_3d = np.array(ann['keypoints_3d']).reshape(-1, 4)
keypoints_3d[..., 3] = keypoints_3d[..., 3] >= 1
bbox = np.array(ann['bbox']).reshape(1, -1)
pose_det_result = {
'image_name': image_name,
'bbox': _xywh2xyxy(bbox),
'keypoints': keypoints,
'keypoints_3d': keypoints_3d
}
pose_det_results.append(pose_det_result)
pose_det_results_list.append(pose_det_results)
else:
print('Stage 1: 2D pose detection.')
pose_det_model = init_pose_model(args.pose_detector_config,
args.pose_detector_checkpoint,
device=args.device.lower())
assert pose_det_model.cfg.model.type == 'TopDown', 'Only "TopDown"' \
'model is supported for the 1st stage (2D pose detection)'
dataset = pose_det_model.cfg.data['test']['type']
dataset_info = pose_det_model.cfg.data['test'].get(
'dataset_info', None)
if dataset_info is None:
warnings.warn(
'Please set `dataset_info` in the config.'
'Check https://github.com/open-mmlab/mmpose/pull/663 '
'for details.', DeprecationWarning)
else:
dataset_info = DatasetInfo(dataset_info)
img_keys = list(coco.imgs.keys())
for i in mmcv.track_iter_progress(range(len(img_keys))):
# get bounding box annotations
image_id = img_keys[i]
image = coco.loadImgs(image_id)[0]
image_name = osp.join(args.img_root, image['file_name'])
ann_ids = coco.getAnnIds(image_id)
# make person results for single image
person_results = []
for ann_id in ann_ids:
person = {}
ann = coco.anns[ann_id]
person['bbox'] = ann['bbox']
person_results.append(person)
pose_det_results, _ = inference_top_down_pose_model(
pose_det_model,
image_name,
person_results,
bbox_thr=None,
format='xywh',
dataset=dataset,
dataset_info=dataset_info,
return_heatmap=False,
outputs=None)
for res in pose_det_results:
res['image_name'] = image_name
pose_det_results_list.append(pose_det_results)
# Second stage: Pose lifting
print('Stage 2: 2D-to-3D pose lifting.')
pose_lift_model = init_pose_model(args.pose_lifter_config,
args.pose_lifter_checkpoint,
device=args.device.lower())
assert pose_lift_model.cfg.model.type == 'PoseLifter', 'Only' \
'"PoseLifter" model is supported for the 2nd stage ' \
'(2D-to-3D lifting)'
dataset = pose_lift_model.cfg.data['test']['type']
dataset_info = pose_lift_model.cfg.data['test'].get('dataset_info', None)
if dataset_info is None:
warnings.warn(
'Please set `dataset_info` in the config.'
'Check https://github.com/open-mmlab/mmpose/pull/663 for details.',
DeprecationWarning)
else:
dataset_info = DatasetInfo(dataset_info)
camera_params = None
if args.camera_param_file is not None:
camera_params = mmcv.load(args.camera_param_file)
for i, pose_det_results in enumerate(
mmcv.track_iter_progress(pose_det_results_list)):
# 2D-to-3D pose lifting
# Note that the pose_det_results are regarded as a single-frame pose
# sequence
pose_lift_results = inference_pose_lifter_model(
pose_lift_model,
pose_results_2d=[pose_det_results],
dataset=dataset,
dataset_info=dataset_info,
with_track_id=False)
image_name = pose_det_results[0]['image_name']
# Pose processing
pose_lift_results_vis = []
for idx, res in enumerate(pose_lift_results):
keypoints_3d = res['keypoints_3d']
# project to world space
if camera_params is not None:
keypoints_3d = _keypoint_camera_to_world(
keypoints_3d,
camera_params=camera_params,
image_name=image_name,
dataset=dataset)
# rebase height (z-axis)
if args.rebase_keypoint_height:
keypoints_3d[..., 2] -= np.min(keypoints_3d[..., 2],
axis=-1,
keepdims=True)
res['keypoints_3d'] = keypoints_3d
# Add title
det_res = pose_det_results[idx]
instance_id = det_res.get('track_id', idx)
res['title'] = f'Prediction ({instance_id})'
pose_lift_results_vis.append(res)
# Add ground truth
if args.show_ground_truth:
if 'keypoints_3d' not in det_res:
print('Fail to show ground truth. Please make sure that'
' the instance annotations from the Json file'
' contain "keypoints_3d".')
else:
gt = res.copy()
gt['keypoints_3d'] = det_res['keypoints_3d']
gt['title'] = f'Ground truth ({instance_id})'
pose_lift_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_lift_model,
result=pose_lift_results_vis,
img=image_name,
dataset_info=dataset_info,
out_file=out_file)
def main():
"""Visualize the demo video with GPU acceleration.
Using full frame to estimate the keypoints.
"""
parser = ArgumentParser()
parser.add_argument('pose_config', help='Config file for pose')
parser.add_argument('pose_checkpoint', help='Checkpoint file for pose')
parser.add_argument('--video-path', type=str, help='Video path')
parser.add_argument(
'--show',
action='store_true',
default=False,
help='whether to show visualizations.')
parser.add_argument(
'--out-video-root',
default='',
help='Root of the output video file. '
'Default not saving the visualization video.')
parser.add_argument(
'--device', default='cuda:0', help='Device used for inference')
parser.add_argument(
'--nvdecode', action='store_true', help='Use NVIDIA decoder')
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_video_root != '')
# 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']
dataset_info = pose_model.cfg.data['test'].get('dataset_info', None)
if dataset_info is None:
warnings.warn(
'Please set `dataset_info` in the config.'
'Check https://github.com/open-mmlab/mmpose/pull/663 for details.',
DeprecationWarning)
else:
dataset_info = DatasetInfo(dataset_info)
pose_model.cfg.data['test']['dataset_info'] = dataset_info
if args.nvdecode:
VideoCapture = ffmpegcv.VideoCaptureNV
else:
VideoCapture = ffmpegcv.VideoCapture
video_origin = VideoCapture(args.video_path)
img_metas = prefetch_img_metas(pose_model.cfg,
(video_origin.width, video_origin.height))
resize_wh = pose_model.cfg.data_cfg['image_size']
video_resize = VideoCapture(
args.video_path,
resize=resize_wh,
resize_keepratio=True,
pix_fmt='rgb24')
if args.out_video_root == '':
save_out_video = False
else:
os.makedirs(args.out_video_root, exist_ok=True)
save_out_video = True
if save_out_video:
fourcc = cv2.VideoWriter_fourcc(*'mp4v')
size = (video_origin.width, video_origin.height)
videoWriter = cv2.VideoWriter(
os.path.join(args.out_video_root,
f'vis_{os.path.basename(args.video_path)}'), fourcc,
video_origin.fps, size)
with torch.no_grad():
for frame_resize, frame_origin in zip(
mmcv.track_iter_progress(video_resize), video_origin):
# test a single image
data = process_img(frame_resize, img_metas, args.device)
pose_results = pose_model(
return_loss=False, return_heatmap=False, **data)
pose_results['keypoints'] = pose_results['preds'][0]
del pose_results['preds']
pose_results = [pose_results]
# show the results
vis_img = vis_pose_result(
pose_model,
frame_origin,
pose_results,
radius=args.radius,
thickness=args.thickness,
dataset=dataset,
dataset_info=dataset_info,
kpt_score_thr=args.kpt_thr,
show=False)
if args.show:
cv2.imshow('Image', vis_img)
if save_out_video:
videoWriter.write(vis_img)
if args.show and cv2.waitKey(1) & 0xFF == ord('q'):
break
video_origin.release()
video_resize.release()
if save_out_video:
videoWriter.release()
if args.show:
cv2.destroyAllWindows()
def main():
"""Visualize the demo video (support both single-frame and multi-frame).
Using mmdet to detect the human.
"""
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_config', help='Config file for pose')
parser.add_argument('pose_checkpoint', help='Checkpoint file for pose')
parser.add_argument('--video-path', type=str, help='Video path')
parser.add_argument(
'--show',
action='store_true',
default=False,
help='whether to show visualizations.')
parser.add_argument(
'--out-video-root',
default='',
help='Root of the output video file. '
'Default not saving the visualization video.')
parser.add_argument(
'--device', default='cuda:0', help='Device used for inference')
parser.add_argument(
'--det-cat-id',
type=int,
default=1,
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, 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')
parser.add_argument(
'--use-multi-frames',
action='store_true',
default=False,
help='whether to use multi frames for inference in the pose'
'estimation stage. Default: False.')
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 pose'
'estimation stage. Default: False.')
assert has_mmdet, 'Please install mmdet to run the demo.'
args = parser.parse_args()
assert args.show or (args.out_video_root != '')
assert args.det_config is not None
assert args.det_checkpoint is not None
print('Initializing model...')
# build the detection model from a config file and a checkpoint file
det_model = init_detector(
args.det_config, args.det_checkpoint, device=args.device.lower())
# 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']
# get datasetinfo
dataset_info = pose_model.cfg.data['test'].get('dataset_info', None)
if dataset_info is None:
warnings.warn(
'Please set `dataset_info` in the config.'
'Check https://github.com/open-mmlab/mmpose/pull/663 for details.',
DeprecationWarning)
else:
dataset_info = DatasetInfo(dataset_info)
# read video
video = mmcv.VideoReader(args.video_path)
assert video.opened, f'Faild to load video file {args.video_path}'
if args.out_video_root == '':
save_out_video = False
else:
os.makedirs(args.out_video_root, exist_ok=True)
save_out_video = True
if save_out_video:
fps = video.fps
size = (video.width, video.height)
fourcc = cv2.VideoWriter_fourcc(*'mp4v')
videoWriter = cv2.VideoWriter(
os.path.join(args.out_video_root,
f'vis_{os.path.basename(args.video_path)}'), fourcc,
fps, size)
# frame index offsets for inference, used in multi-frame inference setting
if args.use_multi_frames:
assert 'frame_indices_test' in pose_model.cfg.data.test.data_cfg
indices = pose_model.cfg.data.test.data_cfg['frame_indices_test']
# whether to return heatmap, optional
return_heatmap = False
# return the output of some desired layers,
# e.g. use ('backbone', ) to return backbone feature
output_layer_names = None
print('Running inference...')
for frame_id, cur_frame in enumerate(mmcv.track_iter_progress(video)):
# get the detection results of current frame
# the resulting box is (x1, y1, x2, y2)
mmdet_results = inference_detector(det_model, cur_frame)
# keep the person class bounding boxes.
person_results = process_mmdet_results(mmdet_results, args.det_cat_id)
if args.use_multi_frames:
frames = collect_multi_frames(video, frame_id, indices,
args.online)
# test a single image, with a list of bboxes.
pose_results, returned_outputs = inference_top_down_pose_model(
pose_model,
frames if args.use_multi_frames else cur_frame,
person_results,
bbox_thr=args.bbox_thr,
format='xyxy',
dataset=dataset,
dataset_info=dataset_info,
return_heatmap=return_heatmap,
outputs=output_layer_names)
# show the results
vis_frame = vis_pose_result(
pose_model,
cur_frame,
pose_results,
dataset=dataset,
dataset_info=dataset_info,
kpt_score_thr=args.kpt_thr,
radius=args.radius,
thickness=args.thickness,
show=False)
if args.show:
cv2.imshow('Frame', vis_frame)
if save_out_video:
videoWriter.write(vis_frame)
if args.show and cv2.waitKey(1) & 0xFF == ord('q'):
break
if save_out_video:
videoWriter.release()
if args.show:
cv2.destroyAllWindows()
def __init__(self,
ann_file,
img_prefix,
data_cfg,
pipeline,
dataset_info=None,
coco_style=True,
test_mode=False):
self.image_info = {}
self.ann_info = {}
self.ann_file = ann_file
self.img_prefix = img_prefix
self.pipeline = pipeline
self.test_mode = test_mode
self.ann_info['image_size'] = np.array(data_cfg['image_size'])
self.ann_info['heatmap_size'] = np.array(data_cfg['heatmap_size'])
self.ann_info['num_joints'] = data_cfg['num_joints']
self.ann_info['inference_channel'] = data_cfg['inference_channel']
self.ann_info['num_output_channels'] = data_cfg['num_output_channels']
self.ann_info['dataset_channel'] = data_cfg['dataset_channel']
self.ann_info['use_different_joint_weights'] = data_cfg.get(
'use_different_joint_weights', False)
if dataset_info is None:
raise ValueError(
'Check https://github.com/open-mmlab/mmpose/pull/663 '
'for details.')
dataset_info = DatasetInfo(dataset_info)
assert self.ann_info['num_joints'] == dataset_info.keypoint_num
self.ann_info['flip_pairs'] = dataset_info.flip_pairs
self.ann_info['flip_index'] = dataset_info.flip_index
self.ann_info['upper_body_ids'] = dataset_info.upper_body_ids
self.ann_info['lower_body_ids'] = dataset_info.lower_body_ids
self.ann_info['joint_weights'] = dataset_info.joint_weights
self.ann_info['skeleton'] = dataset_info.skeleton
self.sigmas = dataset_info.sigmas
self.dataset_name = dataset_info.dataset_name
if coco_style:
self.coco = COCO(ann_file)
if 'categories' in self.coco.dataset:
cats = [
cat['name']
for cat in self.coco.loadCats(self.coco.getCatIds())
]
self.classes = ['__background__'] + cats
self.num_classes = len(self.classes)
self._class_to_ind = dict(
zip(self.classes, range(self.num_classes)))
self._class_to_coco_ind = dict(zip(cats,
self.coco.getCatIds()))
self._coco_ind_to_class_ind = dict(
(self._class_to_coco_ind[cls], self._class_to_ind[cls])
for cls in self.classes[1:])
self.img_ids = self.coco.getImgIds()
self.num_images = len(self.img_ids)
self.id2name, self.name2id = self._get_mapping_id_name(
self.coco.imgs)
self.db = []
self.pipeline = Compose(self.pipeline)
def test_dataset_info():
dataset_info = dict(
dataset_name='zebra',
paper_info=dict(
author='Graving, Jacob M and Chae, Daniel and Naik, Hemal and '
'Li, Liang and Koger, Benjamin and Costelloe, Blair R and '
'Couzin, Iain D',
title='DeepPoseKit, a software toolkit for fast and robust '
'animal pose estimation using deep learning',
container='Elife',
year='2019',
homepage='https://github.com/jgraving/DeepPoseKit-Data',
),
keypoint_info={
0:
dict(name='snout', id=0, color=[255, 255, 255], type='', swap=''),
1:
dict(name='head', id=1, color=[255, 255, 255], type='', swap=''),
2:
dict(name='neck', id=2, color=[255, 255, 255], type='', swap=''),
3:
dict(name='forelegL1',
id=3,
color=[255, 255, 255],
type='',
swap='forelegR1'),
4:
dict(name='forelegR1',
id=4,
color=[255, 255, 255],
type='',
swap='forelegL1'),
5:
dict(name='hindlegL1',
id=5,
color=[255, 255, 255],
type='',
swap='hindlegR1'),
6:
dict(name='hindlegR1',
id=6,
color=[255, 255, 255],
type='',
swap='hindlegL1'),
7:
dict(name='tailbase',
id=7,
color=[255, 255, 255],
type='',
swap=''),
8:
dict(name='tailtip', id=8, color=[255, 255, 255], type='', swap='')
},
skeleton_info={
0: dict(link=('head', 'snout'), id=0, color=[255, 255, 255]),
1: dict(link=('neck', 'head'), id=1, color=[255, 255, 255]),
2: dict(link=('forelegL1', 'neck'), id=2, color=[255, 255, 255]),
3: dict(link=('forelegR1', 'neck'), id=3, color=[255, 255, 255]),
4: dict(link=('hindlegL1', 'tailbase'),
id=4,
color=[255, 255, 255]),
5: dict(link=('hindlegR1', 'tailbase'),
id=5,
color=[255, 255, 255]),
6: dict(link=('tailbase', 'neck'), id=6, color=[255, 255, 255]),
7: dict(link=('tailtip', 'tailbase'), id=7, color=[255, 255, 255])
},
joint_weights=[1.] * 9,
sigmas=[])
dataset_info = DatasetInfo(dataset_info)
assert dataset_info.keypoint_num == len(dataset_info.flip_index)
def main():
"""Visualize the demo images.
Using mmdet to detect the human.
"""
parser = ArgumentParser()
parser.add_argument('pose_config', help='Config file for pose')
parser.add_argument('pose_checkpoint', help='Checkpoint file for pose')
parser.add_argument('--video-path', type=str, help='Video path')
parser.add_argument('--show',
action='store_true',
default=False,
help='whether to show visualizations.')
parser.add_argument('--out-video-root',
default='',
help='Root of the output video file. '
'Default not saving the visualization video.')
parser.add_argument('--device',
default='cuda:0',
help='Device used 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')
assert has_face_det, 'Please install face_recognition to run the demo. '\
'"pip install face_recognition", For more details, '\
'see https://github.com/ageitgey/face_recognition'
args = parser.parse_args()
assert args.show or (args.out_video_root != '')
# 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']
dataset_info = pose_model.cfg.data['test'].get('dataset_info', None)
if dataset_info is None:
warnings.warn(
'Please set `dataset_info` in the config.'
'Check https://github.com/open-mmlab/mmpose/pull/663 for details.',
DeprecationWarning)
else:
dataset_info = DatasetInfo(dataset_info)
cap = cv2.VideoCapture(args.video_path)
assert cap.isOpened(), f'Faild to load video file {args.video_path}'
if args.out_video_root == '':
save_out_video = False
else:
os.makedirs(args.out_video_root, exist_ok=True)
save_out_video = True
if save_out_video:
fps = cap.get(cv2.CAP_PROP_FPS)
size = (int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)),
int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)))
fourcc = cv2.VideoWriter_fourcc(*'mp4v')
videoWriter = cv2.VideoWriter(
os.path.join(args.out_video_root,
f'vis_{os.path.basename(args.video_path)}'), fourcc,
fps, size)
# optional
return_heatmap = False
# e.g. use ('backbone', ) to return backbone feature
output_layer_names = None
while (cap.isOpened()):
flag, img = cap.read()
if not flag:
break
face_det_results = face_recognition.face_locations(
cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
face_results = process_face_det_results(face_det_results)
# test a single image, with a list of bboxes.
pose_results, returned_outputs = inference_top_down_pose_model(
pose_model,
img,
face_results,
bbox_thr=None,
format='xyxy',
dataset=dataset,
dataset_info=dataset_info,
return_heatmap=return_heatmap,
outputs=output_layer_names)
# show the results
vis_img = vis_pose_result(pose_model,
img,
pose_results,
radius=args.radius,
thickness=args.thickness,
dataset=dataset,
dataset_info=dataset_info,
kpt_score_thr=args.kpt_thr,
show=False)
if args.show:
cv2.imshow('Image', vis_img)
if save_out_video:
videoWriter.write(vis_img)
if args.show and cv2.waitKey(1) & 0xFF == ord('q'):
break
cap.release()
if save_out_video:
videoWriter.release()
if args.show:
cv2.destroyAllWindows()
def main():
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_detector_config',
type=str,
default=None,
help='Config file for the 1st stage 2D pose detector')
parser.add_argument(
'pose_detector_checkpoint',
type=str,
default=None,
help='Checkpoint file for the 1st stage 2D pose detector')
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('--video-path',
type=str,
default='',
help='Video path')
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(
'--norm-pose-2d',
action='store_true',
help='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('--show',
action='store_true',
default=False,
help='whether to show visualizations.')
parser.add_argument('--out-video-root',
type=str,
default='vis_results',
help='Root of the output video file. '
'Default not saving the visualization video.')
parser.add_argument('--device',
default='cuda:0',
help='Device for inference')
parser.add_argument('--det-cat-id',
type=int,
default=1,
help='Category id for bounding box detection model')
parser.add_argument('--bbox-thr',
type=float,
default=0.9,
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('--radius',
type=int,
default=8,
help='Keypoint radius for visualization')
parser.add_argument('--thickness',
type=int,
default=2,
help='Link thickness for visualization')
parser.add_argument(
'--smooth',
action='store_true',
help='Apply a temporal filter to smooth the 2D pose estimation '
'results. See also --smooth-filter-cfg.')
parser.add_argument(
'--smooth-filter-cfg',
type=str,
default='configs/_base_/filters/one_euro.py',
help='Config file of the filter to smooth the pose estimation '
'results. See also --smooth.')
parser.add_argument(
'--use-multi-frames',
action='store_true',
default=False,
help='whether to use multi frames for inference in the 2D pose'
'detection stage. Default: False.')
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.')
assert has_mmdet, 'Please install mmdet to run the demo.'
args = parser.parse_args()
assert args.show or (args.out_video_root != '')
assert args.det_config is not None
assert args.det_checkpoint is not None
video = mmcv.VideoReader(args.video_path)
assert video.opened, f'Failed to load video file {args.video_path}'
# First stage: 2D pose detection
print('Stage 1: 2D pose detection.')
print('Initializing model...')
person_det_model = init_detector(args.det_config,
args.det_checkpoint,
device=args.device.lower())
pose_det_model = init_pose_model(args.pose_detector_config,
args.pose_detector_checkpoint,
device=args.device.lower())
assert isinstance(pose_det_model, TopDown), 'Only "TopDown"' \
'model is supported for the 1st stage (2D pose detection)'
# frame index offsets for inference, used in multi-frame inference setting
if args.use_multi_frames:
assert 'frame_indices_test' in pose_det_model.cfg.data.test.data_cfg
indices = pose_det_model.cfg.data.test.data_cfg['frame_indices_test']
pose_det_dataset = pose_det_model.cfg.data['test']['type']
# get datasetinfo
dataset_info = pose_det_model.cfg.data['test'].get('dataset_info', None)
if dataset_info is None:
warnings.warn(
'Please set `dataset_info` in the config.'
'Check https://github.com/open-mmlab/mmpose/pull/663 for details.',
DeprecationWarning)
else:
dataset_info = DatasetInfo(dataset_info)
pose_det_results_list = []
next_id = 0
pose_det_results = []
# whether to return heatmap, optional
return_heatmap = False
# return the output of some desired layers,
# e.g. use ('backbone', ) to return backbone feature
output_layer_names = None
print('Running 2D pose detection inference...')
for frame_id, cur_frame in enumerate(mmcv.track_iter_progress(video)):
pose_det_results_last = pose_det_results
# test a single image, the resulting box is (x1, y1, x2, y2)
mmdet_results = inference_detector(person_det_model, cur_frame)
# keep the person class bounding boxes.
person_det_results = process_mmdet_results(mmdet_results,
args.det_cat_id)
if args.use_multi_frames:
frames = collect_multi_frames(video, frame_id, indices,
args.online)
# make person results for current image
pose_det_results, _ = inference_top_down_pose_model(
pose_det_model,
frames if args.use_multi_frames else cur_frame,
person_det_results,
bbox_thr=args.bbox_thr,
format='xyxy',
dataset=pose_det_dataset,
dataset_info=dataset_info,
return_heatmap=return_heatmap,
outputs=output_layer_names)
# get track id for each person instance
pose_det_results, next_id = get_track_id(
pose_det_results,
pose_det_results_last,
next_id,
use_oks=args.use_oks_tracking,
tracking_thr=args.tracking_thr)
pose_det_results_list.append(copy.deepcopy(pose_det_results))
# Second stage: Pose lifting
print('Stage 2: 2D-to-3D pose lifting.')
print('Initializing model...')
pose_lift_model = init_pose_model(args.pose_lifter_config,
args.pose_lifter_checkpoint,
device=args.device.lower())
assert isinstance(pose_lift_model, PoseLifter), \
'Only "PoseLifter" model is supported for the 2nd stage ' \
'(2D-to-3D lifting)'
pose_lift_dataset = pose_lift_model.cfg.data['test']['type']
if args.out_video_root == '':
save_out_video = False
else:
os.makedirs(args.out_video_root, exist_ok=True)
save_out_video = True
if save_out_video:
fourcc = cv2.VideoWriter_fourcc(*'mp4v')
fps = video.fps
writer = None
# convert keypoint definition
for pose_det_results in pose_det_results_list:
for res in pose_det_results:
keypoints = res['keypoints']
res['keypoints'] = convert_keypoint_definition(
keypoints, pose_det_dataset, pose_lift_dataset)
# load temporal padding config from model.data_cfg
if hasattr(pose_lift_model.cfg, 'test_data_cfg'):
data_cfg = pose_lift_model.cfg.test_data_cfg
else:
data_cfg = pose_lift_model.cfg.data_cfg
# build pose smoother for temporal refinement
if args.smooth:
smoother = Smoother(filter_cfg=args.smooth_filter_cfg,
keypoint_key='keypoints',
keypoint_dim=2)
else:
smoother = None
num_instances = args.num_instances
pose_lift_dataset_info = pose_lift_model.cfg.data['test'].get(
'dataset_info', None)
if pose_lift_dataset_info is None:
warnings.warn(
'Please set `dataset_info` in the config.'
'Check https://github.com/open-mmlab/mmpose/pull/663 for details.',
DeprecationWarning)
else:
pose_lift_dataset_info = DatasetInfo(pose_lift_dataset_info)
print('Running 2D-to-3D pose lifting inference...')
for i, pose_det_results in enumerate(
mmcv.track_iter_progress(pose_det_results_list)):
# extract and pad input pose2d sequence
pose_results_2d = extract_pose_sequence(
pose_det_results_list,
frame_idx=i,
causal=data_cfg.causal,
seq_len=data_cfg.seq_len,
step=data_cfg.seq_frame_interval)
# smooth 2d results
if smoother:
pose_results_2d = smoother.smooth(pose_results_2d)
# 2D-to-3D pose lifting
pose_lift_results = inference_pose_lifter_model(
pose_lift_model,
pose_results_2d=pose_results_2d,
dataset=pose_lift_dataset,
dataset_info=pose_lift_dataset_info,
with_track_id=True,
image_size=video.resolution,
norm_pose_2d=args.norm_pose_2d)
# Pose processing
pose_lift_results_vis = []
for idx, res in enumerate(pose_lift_results):
keypoints_3d = res['keypoints_3d']
# exchange y,z-axis, and then reverse the direction of x,z-axis
keypoints_3d = keypoints_3d[..., [0, 2, 1]]
keypoints_3d[..., 0] = -keypoints_3d[..., 0]
keypoints_3d[..., 2] = -keypoints_3d[..., 2]
# rebase height (z-axis)
if args.rebase_keypoint_height:
keypoints_3d[..., 2] -= np.min(keypoints_3d[..., 2],
axis=-1,
keepdims=True)
res['keypoints_3d'] = keypoints_3d
# add title
det_res = pose_det_results[idx]
instance_id = det_res['track_id']
res['title'] = f'Prediction ({instance_id})'
# only visualize the target frame
res['keypoints'] = det_res['keypoints']
res['bbox'] = det_res['bbox']
res['track_id'] = instance_id
pose_lift_results_vis.append(res)
# Visualization
if num_instances < 0:
num_instances = len(pose_lift_results_vis)
img_vis = vis_3d_pose_result(pose_lift_model,
result=pose_lift_results_vis,
img=video[i],
dataset=pose_lift_dataset,
dataset_info=pose_lift_dataset_info,
out_file=None,
radius=args.radius,
thickness=args.thickness,
num_instances=num_instances,
show=args.show)
if save_out_video:
if writer is None:
writer = cv2.VideoWriter(
osp.join(args.out_video_root,
f'vis_{osp.basename(args.video_path)}'), fourcc,
fps, (img_vis.shape[1], img_vis.shape[0]))
writer.write(img_vis)
if save_out_video:
writer.release()
