def get_pose_estimation_prediction(pose_model, image, center, scale): rotation = 0 # pose estimation transformation trans = get_affine_transform(center, scale, rotation, cfg.MODEL.IMAGE_SIZE) model_input = cv2.warpAffine( image, trans, (int(cfg.MODEL.IMAGE_SIZE[0]), int(cfg.MODEL.IMAGE_SIZE[1])), flags=cv2.INTER_LINEAR) transform = transforms.Compose([ transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), ]) # pose estimation inference model_input = transform(model_input).unsqueeze(0) # switch to evaluate mode pose_model.eval() with torch.no_grad(): # compute output heatmap output = pose_model(model_input.to(device)) preds, _ = get_final_preds( cfg, output.clone().cpu().numpy(), np.asarray([center]), np.asarray([scale])) return preds
def main(): args = parse_args() update_config(cfg, args) logger, final_output_dir, tb_log_dir = create_logger( cfg, args.cfg, 'valid') logger.info(pprint.pformat(args)) logger.info(cfg) # cudnn related setting cudnn.benchmark = cfg.CUDNN.BENCHMARK torch.backends.cudnn.deterministic = cfg.CUDNN.DETERMINISTIC torch.backends.cudnn.enabled = cfg.CUDNN.ENABLED model = eval('models.' + cfg.MODEL.NAME + '.get_pose_net')(cfg, is_train=False) if cfg.TEST.MODEL_FILE: logger.info('=> loading model from {}'.format(cfg.TEST.MODEL_FILE)) model.load_state_dict(torch.load(cfg.TEST.MODEL_FILE, map_location={'cuda:0': 'cpu'}), strict=False) else: model_state_file = os.path.join(final_output_dir, 'final_state.pth') logger.info('=> loading model from {}'.format(model_state_file)) model.load_state_dict(torch.load(model_state_file)) model = torch.nn.DataParallel(model, device_ids=[0]) #.cuda() # define loss function (criterion) and optimizer criterion = JointsMSELoss( use_target_weight=cfg.LOSS.USE_TARGET_WEIGHT) #.cuda() # Loading an image image_file = args.img_file data_numpy = cv2.imread(image_file, cv2.IMREAD_COLOR | cv2.IMREAD_IGNORE_ORIENTATION) if data_numpy is None: logger.error('=> fail to read {}'.format(image_file)) raise ValueError('=> fail to read {}'.format(image_file)) # object detection box box = [0, 0, data_numpy.shape[0], data_numpy.shape[1]] c, s = _box2cs(box, data_numpy.shape[1], data_numpy.shape[0]) r = 0 trans = get_affine_transform(c, s, r, cfg.MODEL.IMAGE_SIZE) input = cv2.warpAffine( data_numpy, trans, (int(cfg.MODEL.IMAGE_SIZE[0]), int(cfg.MODEL.IMAGE_SIZE[1])), flags=cv2.INTER_LINEAR) transform = transforms.Compose([ transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), ]) input = transform(input).unsqueeze(0) # switch to evaluate mode model.eval() with torch.no_grad(): # compute output heatmap output = model(input) preds, maxvals = get_final_preds(cfg, output.clone().cpu().numpy(), np.asarray([c]), np.asarray([s])) image = data_numpy.copy() for mat in preds[0]: x, y = int(mat[0]), int(mat[1]) cv2.circle(image, (x, y), 2, (255, 0, 0), 2) # vis result cv2.imwrite("test_h36m.jpg", image) cv2.imshow('res', image) cv2.waitKey(10000)