def inference_bottom_up_pose_model(model,
img_or_path,
pose_nms_thr=0.9,
return_heatmap=False,
outputs=None):
"""Inference a single image.
num_people: P
num_keypoints: K
bbox height: H
bbox width: W
Args:
model (nn.Module): The loaded pose model.
img_or_path (str| np.ndarray): Image filename or loaded image.
pose_nms_thr (float): retain oks overlap < pose_nms_thr, default: 0.9.
return_heatmap (bool) : Flag to return heatmap, default: False.
outputs (list(str) | tuple(str)) : Names of layers whose outputs
need to be returned, default: None.
Returns:
list[ndarray]: The predicted pose info.
The length of the list is the number of people (P).
Each item in the list is a ndarray, containing each person's
pose (ndarray[Kx3]): x, y, score.
list[dict[np.ndarray[N, K, H, W] | torch.tensor[N, K, H, W]]]:
Output feature maps from layers specified in `outputs`.
Includes 'heatmap' if `return_heatmap` is True.
"""
pose_results = []
returned_outputs = []
cfg = model.cfg
device = next(model.parameters()).device
# build the data pipeline
channel_order = cfg.test_pipeline[0].get('channel_order', 'rgb')
test_pipeline = [LoadImage(channel_order=channel_order)
] + cfg.test_pipeline[1:]
test_pipeline = Compose(test_pipeline)
# prepare data
data = {
'img_or_path': img_or_path,
'dataset': 'coco',
'ann_info': {
'image_size':
cfg.data_cfg['image_size'],
'num_joints':
cfg.data_cfg['num_joints'],
'flip_index':
[0, 2, 1, 4, 3, 6, 5, 8, 7, 10, 9, 12, 11, 14, 13, 16, 15],
}
}
data = test_pipeline(data)
data = collate([data], samples_per_gpu=1)
if next(model.parameters()).is_cuda:
# scatter to specified GPU
data = scatter(data, [device])[0]
else:
# just get the actual data from DataContainer
data['img_metas'] = data['img_metas'].data[0]
with OutputHook(model, outputs=outputs, as_tensor=False) as h:
# forward the model
with torch.no_grad():
result = model(img=data['img'],
img_metas=data['img_metas'],
return_loss=False,
return_heatmap=return_heatmap)
if return_heatmap:
h.layer_outputs['heatmap'] = result['output_heatmap']
returned_outputs.append(h.layer_outputs)
for idx, pred in enumerate(result['preds']):
area = (np.max(pred[:, 0]) - np.min(pred[:, 0])) * (
np.max(pred[:, 1]) - np.min(pred[:, 1]))
pose_results.append({
'keypoints': pred[:, :3],
'score': result['scores'][idx],
'area': area,
})
# pose nms
keep = oks_nms(pose_results, pose_nms_thr, sigmas=None)
pose_results = [pose_results[_keep] for _keep in keep]
return pose_results, returned_outputs
def inference_bottom_up_pose_model(model,
img_or_path,
dataset='BottomUpCocoDataset',
dataset_info=None,
pose_nms_thr=0.9,
return_heatmap=False,
outputs=None):
"""Inference a single image with a bottom-up pose model.
Note:
- num_people: P
- num_keypoints: K
- bbox height: H
- bbox width: W
Args:
model (nn.Module): The loaded pose model.
img_or_path (str| np.ndarray): Image filename or loaded image.
dataset (str): Dataset name, e.g. 'BottomUpCocoDataset'.
It is deprecated. Please use dataset_info instead.
dataset_info (DatasetInfo): A class containing all dataset info.
pose_nms_thr (float): retain oks overlap < pose_nms_thr, default: 0.9.
return_heatmap (bool) : Flag to return heatmap, default: False.
outputs (list(str) | tuple(str)) : Names of layers whose outputs
need to be returned, default: None.
Returns:
tuple:
- pose_results (list[np.ndarray]): The predicted pose info. \
The length of the list is the number of people (P). \
Each item in the list is a ndarray, containing each \
person's pose (np.ndarray[Kx3]): x, y, score.
- returned_outputs (list[dict[np.ndarray[N, K, H, W] | \
torch.Tensor[N, K, H, W]]]): \
Output feature maps from layers specified in `outputs`. \
Includes 'heatmap' if `return_heatmap` is True.
"""
# get dataset info
if (dataset_info is None and hasattr(model, 'cfg')
and 'dataset_info' in model.cfg):
dataset_info = DatasetInfo(model.cfg.dataset_info)
if dataset_info is not None:
dataset_name = dataset_info.dataset_name
flip_index = dataset_info.flip_index
sigmas = getattr(dataset_info, 'sigmas', None)
else:
warnings.warn(
'dataset is deprecated.'
'Please set `dataset_info` in the config.'
'Check https://github.com/open-mmlab/mmpose/pull/663 for details.',
DeprecationWarning)
assert (dataset == 'BottomUpCocoDataset')
dataset_name = dataset
flip_index = [0, 2, 1, 4, 3, 6, 5, 8, 7, 10, 9, 12, 11, 14, 13, 16, 15]
sigmas = None
pose_results = []
returned_outputs = []
cfg = model.cfg
device = next(model.parameters()).device
if device.type == 'cpu':
device = -1
# build the data pipeline
test_pipeline = Compose(cfg.test_pipeline)
_pipeline_gpu_speedup(test_pipeline, next(model.parameters()).device)
# prepare data
data = {
'dataset': dataset_name,
'ann_info': {
'image_size': np.array(cfg.data_cfg['image_size']),
'num_joints': cfg.data_cfg['num_joints'],
'flip_index': flip_index,
}
}
if isinstance(img_or_path, np.ndarray):
data['img'] = img_or_path
else:
data['image_file'] = img_or_path
data = test_pipeline(data)
data = collate([data], samples_per_gpu=1)
data = scatter(data, [device])[0]
with OutputHook(model, outputs=outputs, as_tensor=False) as h:
# forward the model
with torch.no_grad():
result = model(img=data['img'],
img_metas=data['img_metas'],
return_loss=False,
return_heatmap=return_heatmap)
if return_heatmap:
h.layer_outputs['heatmap'] = result['output_heatmap']
returned_outputs.append(h.layer_outputs)
for idx, pred in enumerate(result['preds']):
area = (np.max(pred[:, 0]) - np.min(pred[:, 0])) * (
np.max(pred[:, 1]) - np.min(pred[:, 1]))
pose_results.append({
'keypoints': pred[:, :3],
'score': result['scores'][idx],
'area': area,
})
# pose nms
score_per_joint = cfg.model.test_cfg.get('score_per_joint', False)
keep = oks_nms(pose_results,
pose_nms_thr,
sigmas,
score_per_joint=score_per_joint)
pose_results = [pose_results[_keep] for _keep in keep]
return pose_results, returned_outputs