Spaces:
Sleeping
Sleeping
| # Copyright (c) OpenMMLab. All rights reserved. | |
| import logging | |
| from argparse import ArgumentParser | |
| from mmcv.image import imread | |
| from mmengine.logging import print_log | |
| from mmpose.apis import inference_topdown, init_model | |
| from mmpose.registry import VISUALIZERS | |
| from mmpose.structures import merge_data_samples | |
| def parse_args(): | |
| parser = ArgumentParser() | |
| parser.add_argument('img', help='Image file') | |
| parser.add_argument('config', help='Config file') | |
| parser.add_argument('checkpoint', help='Checkpoint file') | |
| parser.add_argument('--out-file', default=None, help='Path to output file') | |
| parser.add_argument( | |
| '--device', default='cuda:0', help='Device used for inference') | |
| parser.add_argument( | |
| '--draw-heatmap', | |
| action='store_true', | |
| help='Visualize the predicted heatmap') | |
| parser.add_argument( | |
| '--show-kpt-idx', | |
| action='store_true', | |
| default=False, | |
| help='Whether to show the index of keypoints') | |
| parser.add_argument( | |
| '--skeleton-style', | |
| default='mmpose', | |
| type=str, | |
| choices=['mmpose', 'openpose'], | |
| help='Skeleton style selection') | |
| parser.add_argument( | |
| '--kpt-thr', | |
| type=float, | |
| default=0.3, | |
| help='Visualizing keypoint thresholds') | |
| parser.add_argument( | |
| '--radius', | |
| type=int, | |
| default=3, | |
| help='Keypoint radius for visualization') | |
| parser.add_argument( | |
| '--thickness', | |
| type=int, | |
| default=1, | |
| help='Link thickness for visualization') | |
| parser.add_argument( | |
| '--alpha', type=float, default=0.8, help='The transparency of bboxes') | |
| parser.add_argument( | |
| '--show', | |
| action='store_true', | |
| default=False, | |
| help='whether to show img') | |
| args = parser.parse_args() | |
| return args | |
| def main(): | |
| args = parse_args() | |
| # build the model from a config file and a checkpoint file | |
| if args.draw_heatmap: | |
| cfg_options = dict(model=dict(test_cfg=dict(output_heatmaps=True))) | |
| else: | |
| cfg_options = None | |
| model = init_model( | |
| args.config, | |
| args.checkpoint, | |
| device=args.device, | |
| cfg_options=cfg_options) | |
| # init visualizer | |
| model.cfg.visualizer.radius = args.radius | |
| model.cfg.visualizer.alpha = args.alpha | |
| model.cfg.visualizer.line_width = args.thickness | |
| visualizer = VISUALIZERS.build(model.cfg.visualizer) | |
| visualizer.set_dataset_meta( | |
| model.dataset_meta, skeleton_style=args.skeleton_style) | |
| # inference a single image | |
| batch_results = inference_topdown(model, args.img) | |
| results = merge_data_samples(batch_results) | |
| # show the results | |
| img = imread(args.img, channel_order='rgb') | |
| visualizer.add_datasample( | |
| 'result', | |
| img, | |
| data_sample=results, | |
| draw_gt=False, | |
| draw_bbox=True, | |
| kpt_thr=args.kpt_thr, | |
| draw_heatmap=args.draw_heatmap, | |
| show_kpt_idx=args.show_kpt_idx, | |
| skeleton_style=args.skeleton_style, | |
| show=args.show, | |
| out_file=args.out_file) | |
| if args.out_file is not None: | |
| print_log( | |
| f'the output image has been saved at {args.out_file}', | |
| logger='current', | |
| level=logging.INFO) | |
| if __name__ == '__main__': | |
| main() | |