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()