def worker(video_file, index, detection_cfg, skeleton_cfg, skeleon_data_cfg,
device, result_queue):
os.environ["CUDA_VISIBLE_DEVICES"] = str(device)
video_frames = mmcv.VideoReader(video_file)
# load model
detection_model_file = detection_cfg.model_cfg
detection_checkpoint_file = get_mmskeleton_url(
detection_cfg.checkpoint_file)
detection_model = init_detector(detection_model_file,
detection_checkpoint_file,
device='cpu')
skeleton_model_file = skeleton_cfg.model_cfg
skeletion_checkpoint_file = skeleton_cfg.checkpoint_file
skeleton_model = init_twodimestimator(skeleton_model_file,
skeletion_checkpoint_file,
device='cpu')
detection_model = detection_model.cuda()
skeleton_model = skeleton_model.cuda()
for idx in index:
skeleton_result = dict()
image = video_frames[idx]
draw_image = image.copy()
bbox_result = inference_detector(detection_model, image)
person_bbox, labels = VideoDemo.bbox_filter(bbox_result,
detection_cfg.bbox_thre)
if len(person_bbox) > 0:
person, meta = VideoDemo.skeleton_preprocess(
image[:, :, ::-1], person_bbox, skeleon_data_cfg)
preds, maxvals = inference_twodimestimator(skeleton_model,
person.cuda(), meta,
True)
results = VideoDemo.skeleton_postprocess(preds, maxvals, meta)
if skeleon_data_cfg.save_video:
file = os.path.join(skeleon_data_cfg.img_dir,
'{}.png'.format(idx))
mmcv.imshow_det_bboxes(draw_image,
person_bbox,
labels,
detection_model.CLASSES,
score_thr=detection_cfg.bbox_thre,
show=False,
wait_time=0)
save(image, draw_image, results, file)
else:
preds, maxvals = None, None
if skeleon_data_cfg.save_video:
file = os.path.join(skeleon_data_cfg.img_dir,
'{}.png'.format(idx))
mmcv.imwrite(image, file)
skeleton_result['frame_index'] = idx
skeleton_result['position_preds'] = preds
skeleton_result['position_maxvals'] = maxvals
result_queue.put(skeleton_result)
def init_pose_estimator(detection_cfg, estimation_cfg, device=None):
detection_model_file = detection_cfg.model_cfg
detection_checkpoint_file = get_mmskeleton_url(
detection_cfg.checkpoint_file)
detection_model = mmdet.apis.init_detector(detection_model_file,
detection_checkpoint_file,
device='cpu')
skeleton_model_file = estimation_cfg.model_cfg
skeletion_checkpoint_file = estimation_cfg.checkpoint_file
skeleton_model = init_twodimestimator(skeleton_model_file,
skeletion_checkpoint_file,
device='cpu')
if device is not None:
os.environ["CUDA_VISIBLE_DEVICES"] = str(device)
detection_model = detection_model.cuda()
skeleton_model = skeleton_model.cuda()
pose_estimator = (detection_model, skeleton_model, detection_cfg,
estimation_cfg)
return pose_estimator