Spaces:
Running
on
A10G
Running
on
A10G
# Copyright (c) OpenMMLab. All rights reserved. | |
from abc import ABCMeta, abstractmethod | |
import numpy as np | |
import torch.nn as nn | |
from mmpose.core.evaluation.top_down_eval import keypoints_from_heatmaps | |
class TopdownHeatmapBaseHead(nn.Module): | |
"""Base class for top-down heatmap heads. | |
All top-down heatmap heads should subclass it. | |
All subclass should overwrite: | |
Methods:`get_loss`, supporting to calculate loss. | |
Methods:`get_accuracy`, supporting to calculate accuracy. | |
Methods:`forward`, supporting to forward model. | |
Methods:`inference_model`, supporting to inference model. | |
""" | |
__metaclass__ = ABCMeta | |
def get_loss(self, **kwargs): | |
"""Gets the loss.""" | |
def get_accuracy(self, **kwargs): | |
"""Gets the accuracy.""" | |
def forward(self, **kwargs): | |
"""Forward function.""" | |
def inference_model(self, **kwargs): | |
"""Inference function.""" | |
def decode(self, img_metas, output, **kwargs): | |
"""Decode keypoints from heatmaps. | |
Args: | |
img_metas (list(dict)): Information about data augmentation | |
By default this includes: | |
- "image_file: path to the image file | |
- "center": center of the bbox | |
- "scale": scale of the bbox | |
- "rotation": rotation of the bbox | |
- "bbox_score": score of bbox | |
output (np.ndarray[N, K, H, W]): model predicted heatmaps. | |
""" | |
batch_size = len(img_metas) | |
if 'bbox_id' in img_metas[0]: | |
bbox_ids = [] | |
else: | |
bbox_ids = None | |
c = np.zeros((batch_size, 2), dtype=np.float32) | |
s = np.zeros((batch_size, 2), dtype=np.float32) | |
image_paths = [] | |
score = np.ones(batch_size) | |
for i in range(batch_size): | |
c[i, :] = img_metas[i]['center'] | |
s[i, :] = img_metas[i]['scale'] | |
image_paths.append(img_metas[i]['image_file']) | |
if 'bbox_score' in img_metas[i]: | |
score[i] = np.array(img_metas[i]['bbox_score']).reshape(-1) | |
if bbox_ids is not None: | |
bbox_ids.append(img_metas[i]['bbox_id']) | |
preds, maxvals = keypoints_from_heatmaps( | |
output, | |
c, | |
s, | |
unbiased=self.test_cfg.get('unbiased_decoding', False), | |
post_process=self.test_cfg.get('post_process', 'default'), | |
kernel=self.test_cfg.get('modulate_kernel', 11), | |
valid_radius_factor=self.test_cfg.get('valid_radius_factor', | |
0.0546875), | |
use_udp=self.test_cfg.get('use_udp', False), | |
target_type=self.test_cfg.get('target_type', 'GaussianHeatmap')) | |
all_preds = np.zeros((batch_size, preds.shape[1], 3), dtype=np.float32) | |
all_boxes = np.zeros((batch_size, 6), dtype=np.float32) | |
all_preds[:, :, 0:2] = preds[:, :, 0:2] | |
all_preds[:, :, 2:3] = maxvals | |
all_boxes[:, 0:2] = c[:, 0:2] | |
all_boxes[:, 2:4] = s[:, 0:2] | |
all_boxes[:, 4] = np.prod(s * 200.0, axis=1) | |
all_boxes[:, 5] = score | |
result = {} | |
result['preds'] = all_preds | |
result['boxes'] = all_boxes | |
result['image_paths'] = image_paths | |
result['bbox_ids'] = bbox_ids | |
return result | |
def _get_deconv_cfg(deconv_kernel): | |
"""Get configurations for deconv layers.""" | |
if deconv_kernel == 4: | |
padding = 1 | |
output_padding = 0 | |
elif deconv_kernel == 3: | |
padding = 1 | |
output_padding = 1 | |
elif deconv_kernel == 2: | |
padding = 0 | |
output_padding = 0 | |
else: | |
raise ValueError(f'Not supported num_kernels ({deconv_kernel}).') | |
return deconv_kernel, padding, output_padding | |