def inference( detection_cfg, skeleton_cfg, dataset_cfg, gpus=1, worker_per_gpu=1, ): # get frame num video_file = dataset_cfg.video_file video_name = video_file.strip('/n').split('/')[-1] video_frames = mmcv.VideoReader(video_file) num_frames = len(video_frames) del video_frames data_cfg = skeleton_cfg.data_cfg if data_cfg.save_video: data_cfg.img_dir = os.path.join(data_cfg.save_dir, '{}.img'.format(video_name)) if os.path.exists(data_cfg.img_dir): import shutil shutil.rmtree(data_cfg.img_dir) os.makedirs(data_cfg.img_dir) # cache model checkpoints cache_checkpoint(detection_cfg.checkpoint_file) cache_checkpoint(skeleton_cfg.checkpoint_file) # multiprocess settings context = mp.get_context('spawn') result_queue = context.Queue(num_frames) procs = [] for w in range(gpus * worker_per_gpu): shred_list = list(range(w, num_frames, gpus * worker_per_gpu)) p = context.Process(target=worker, args=(video_file, shred_list, detection_cfg, skeleton_cfg, data_cfg, w % gpus, result_queue)) p.start() procs.append(p) all_result = [] print('\nPose estimation start:') prog_bar = ProgressBar(num_frames) for i in range(num_frames): t = result_queue.get() all_result.append(t) prog_bar.update() for p in procs: p.join() if len(all_result) == num_frames and data_cfg.save_video: print('\n\nGenerate video:') video_path = os.path.join(data_cfg.save_dir, video_name) mmcv.frames2video(data_cfg.img_dir, video_path, filename_tmpl='{:01d}.png') print('Video was saved to {}'.format(video_path)) import IPython IPython.embed()
def build(detection_cfg, estimation_cfg, tracker_cfg, video_dir, out_dir, gpus=1, worker_per_gpu=1, video_max_length=10000, category_annotation=None): cache_checkpoint(detection_cfg.checkpoint_file) cache_checkpoint(estimation_cfg.checkpoint_file) if tracker_cfg is not None: raise NotImplementedError if not os.path.isdir(out_dir): os.makedirs(out_dir) if category_annotation is None: video_categories = dict() else: with open(category_annotation) as f: video_categories = json.load(f)['annotations'] inputs = Manager().Queue(video_max_length) results = Manager().Queue(video_max_length) num_worker = gpus * worker_per_gpu procs = [] for i in range(num_worker): p = Process(target=worker, args=(inputs, results, i % gpus, detection_cfg, estimation_cfg)) procs.append(p) p.start() video_file_list = os.listdir(video_dir) prog_bar = ProgressBar(len(video_file_list)) for video_file in video_file_list: #print('video_file:',video_file) reader = mmcv.VideoReader(os.path.join(video_dir, video_file)) video_frames = reader[:video_max_length] annotations = [] num_keypoints = -1 count_frame = 0 for i, image in enumerate(video_frames): #print('input_i:',i) if image is None: continue count_frame += 1 inputs.put((i, image)) #print('\ncount_frame:',count_frame) #print('len_frame',len(video_frames)) for i in range(count_frame): #print('frame:',i) t = results.get() #print('t:',t) if not t['has_return']: continue num_person = len(t['joint_preds']) assert len(t['person_bbox']) == num_person for j in range(num_person): keypoints = [[p[0], p[1], round(s[0], 2)] for p, s in zip( t['joint_preds'][j].round().astype(int).tolist(), t['joint_scores'][j].tolist())] num_keypoints = len(keypoints) person_info = dict(person_bbox=t['person_bbox'] [j].round().astype(int).tolist(), frame_index=t['frame_index'], id=j, person_id=None, keypoints=keypoints) annotations.append(person_info) # output results annotations = sorted(annotations, key=lambda x: x['frame_index']) category_id = video_categories[video_file][ 'category_id'] if video_file in video_categories else -1 info = dict(video_name=video_file, resolution=reader.resolution, num_frame=len(video_frames), num_keypoints=num_keypoints, keypoint_channels=['x', 'y', 'score'], version='1.0') video_info = dict(info=info, category_id=category_id, annotations=annotations) with open(os.path.join(out_dir, video_file + '.json'), 'w') as f: json.dump(video_info, f) prog_bar.update() # send end signals for p in procs: inputs.put((-1, None)) # wait to finish for p in procs: p.join() print('\nBuild skeleton dataset to {}.'.format(out_dir)) return video_info
def inference(detection_cfg, estimation_cfg, video_file, gpus=1, worker_per_gpu=1, save_dir=None): video_frames = mmcv.VideoReader(video_file) all_result = [] print('\nPose estimation:') # case for single process if gpus == 1 and worker_per_gpu == 1: model = init_pose_estimator(detection_cfg, estimation_cfg, device=0) prog_bar = ProgressBar(len(video_frames)) for i, image in enumerate(video_frames): res = inference_pose_estimator(model, image) res['frame_index'] = i if save_dir is not None: res['render_image'] = render(image, res['joint_preds'], res['person_bbox'], detection_cfg.bbox_thre) all_result.append(res) prog_bar.update() # case for multi-process else: cache_checkpoint(detection_cfg.checkpoint_file) cache_checkpoint(estimation_cfg.checkpoint_file) num_worker = gpus * worker_per_gpu procs = [] inputs = Manager().Queue(len(video_frames)) results = Manager().Queue(len(video_frames)) for i, image in enumerate(video_frames): inputs.put((i, image)) for i in range(num_worker): p = Process(target=worker, args=(inputs, results, i % gpus, detection_cfg, estimation_cfg, save_dir is not None)) procs.append(p) p.start() for i in range(len(video_frames)): t = results.get() all_result.append(t) if 'prog_bar' not in locals(): prog_bar = ProgressBar(len(video_frames)) prog_bar.update() for p in procs: p.join() # sort results all_result = sorted(all_result, key=lambda x: x['frame_index']) # generate video if (len(all_result) == len(video_frames)) and (save_dir is not None): print('\n\nGenerate video:') video_name = video_file.strip('/n').split('/')[-1] video_path = os.path.join(save_dir, video_name) vwriter = cv2.VideoWriter(video_path, mmcv.video.io.VideoWriter_fourcc(*('mp4v')), video_frames.fps, video_frames.resolution) prog_bar = ProgressBar(len(video_frames)) for r in all_result: vwriter.write(r['render_image']) prog_bar.update() vwriter.release() print('\nVideo was saved to {}'.format(video_path)) return all_result
def build(inputs, detection_cfg, estimation_cfg, tracker_cfg, video_dir, gpus=1, video_max_length=10000, category_annotation=None): print('data build start') cache_checkpoint(detection_cfg.checkpoint_file) cache_checkpoint(estimation_cfg.checkpoint_file) if category_annotation is None: video_categories = dict() else: with open(category_annotation) as f: video_categories = json.load(f)['annotations'] if tracker_cfg is not None: raise NotImplementedError pose_estimators = init_pose_estimator(detection_cfg, estimation_cfg, device=0) video_file_list = [] get_all_file(video_dir, video_file_list) prog_bar = ProgressBar(len(video_file_list)) for video_path in video_file_list: video_file = os.path.basename(video_path) reader = mmcv.VideoReader(video_path) video_frames = reader[:video_max_length] annotations = [] num_keypoints = -1 for i, image in enumerate(video_frames): res = inference_pose_estimator(pose_estimators, image) res['frame_index'] = i if not res['has_return']: continue num_person = len(res['joint_preds']) assert len(res['person_bbox']) == num_person for j in range(num_person): keypoints = [[p[0], p[1], round(s[0], 2)] for p, s in zip( res['joint_preds'][j].round().astype(int).tolist(), res['joint_scores'][j].tolist())] num_keypoints = len(keypoints) person_info = dict(person_bbox=res['person_bbox'] [j].round().astype(int).tolist(), frame_index=res['frame_index'], id=j, person_id=None, keypoints=keypoints) annotations.append(person_info) annotations = sorted(annotations, key=lambda x: x['frame_index']) category_id = video_categories[video_file][ 'category_id'] if video_file in video_categories else -1 info = dict(video_name=video_file, resolution=reader.resolution, num_frame=len(video_frames), num_keypoints=num_keypoints, keypoint_channels=['x', 'y', 'score'], version='1.0') video_info = dict(info=info, category_id=category_id, annotations=annotations) inputs.put(video_info) prog_bar.update()