Example #1
0
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)
Example #2
0
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