mmpose.models.heads.heatmap_heads.cpm_head 源代码

# Copyright (c) OpenMMLab. All rights reserved.
from typing import List, Optional, Sequence, Union

import torch
from mmcv.cnn import build_conv_layer, build_upsample_layer
from mmengine.structures import PixelData
from torch import Tensor, nn

from mmpose.evaluation.functional import pose_pck_accuracy
from mmpose.models.utils.tta import flip_heatmaps
from mmpose.registry import KEYPOINT_CODECS, MODELS
from mmpose.utils.tensor_utils import to_numpy
from mmpose.utils.typing import (Features, MultiConfig, OptConfigType,
                                 OptSampleList, Predictions)
from ..base_head import BaseHead

OptIntSeq = Optional[Sequence[int]]

[文档]@MODELS.register_module() class CPMHead(BaseHead): """Multi-stage heatmap head introduced in `Convolutional Pose Machines`_ by Wei et al (2016) and used by `Stacked Hourglass Networks`_ by Newell et al (2016). The head consists of multiple branches, each of which has some deconv layers and a simple conv2d layer. Args: in_channels (int | Sequence[int]): Number of channels in the input feature maps. out_channels (int): Number of channels in the output heatmaps. num_stages (int): Number of stages. deconv_out_channels (Sequence[int], optional): The output channel number of each deconv layer. Defaults to ``(256, 256, 256)`` deconv_kernel_sizes (Sequence[int | tuple], optional): The kernel size of each deconv layer. Each element should be either an integer for both height and width dimensions, or a tuple of two integers for the height and the width dimension respectively. Defaults to ``(4, 4, 4)`` final_layer (dict): Arguments of the final Conv2d layer. Defaults to ``dict(kernel_size=1)`` loss (Config | List[Config]): Config of the keypoint loss of different stages. Defaults to use :class:`KeypointMSELoss`. decoder (Config, optional): The decoder config that controls decoding keypoint coordinates from the network output. Defaults to ``None`` init_cfg (Config, optional): Config to control the initialization. See :attr:`default_init_cfg` for default settings .. _`Convolutional Pose Machines`: .. _`Stacked Hourglass Networks`: """ _version = 2 def __init__(self, in_channels: Union[int, Sequence[int]], out_channels: int, num_stages: int, deconv_out_channels: OptIntSeq = None, deconv_kernel_sizes: OptIntSeq = None, final_layer: dict = dict(kernel_size=1), loss: MultiConfig = dict( type='KeypointMSELoss', use_target_weight=True), decoder: OptConfigType = None, init_cfg: OptConfigType = None): if init_cfg is None: init_cfg = self.default_init_cfg super().__init__(init_cfg) self.num_stages = num_stages self.in_channels = in_channels self.out_channels = out_channels if isinstance(loss, list): if len(loss) != num_stages: raise ValueError( f'The length of loss_module({len(loss)}) did not match ' f'`num_stages`({num_stages})') self.loss_module = nn.ModuleList( for _loss in loss) else: self.loss_module = if decoder is not None: self.decoder = else: self.decoder = None # build multi-stage deconv layers self.multi_deconv_layers = nn.ModuleList([]) if deconv_out_channels: if deconv_kernel_sizes is None or len(deconv_out_channels) != len( deconv_kernel_sizes): raise ValueError( '"deconv_out_channels" and "deconv_kernel_sizes" should ' 'be integer sequences with the same length. Got ' f'mismatched lengths {deconv_out_channels} and ' f'{deconv_kernel_sizes}') for _ in range(self.num_stages): deconv_layers = self._make_deconv_layers( in_channels=in_channels, layer_out_channels=deconv_out_channels, layer_kernel_sizes=deconv_kernel_sizes, ) self.multi_deconv_layers.append(deconv_layers) in_channels = deconv_out_channels[-1] else: for _ in range(self.num_stages): self.multi_deconv_layers.append(nn.Identity()) # build multi-stage final layers self.multi_final_layers = nn.ModuleList([]) if final_layer is not None: cfg = dict( type='Conv2d', in_channels=in_channels, out_channels=out_channels, kernel_size=1) cfg.update(final_layer) for _ in range(self.num_stages): self.multi_final_layers.append(build_conv_layer(cfg)) else: for _ in range(self.num_stages): self.multi_final_layers.append(nn.Identity()) @property def default_init_cfg(self): init_cfg = [ dict( type='Normal', layer=['Conv2d', 'ConvTranspose2d'], std=0.001), dict(type='Constant', layer='BatchNorm2d', val=1) ] return init_cfg def _make_deconv_layers(self, in_channels: int, layer_out_channels: Sequence[int], layer_kernel_sizes: Sequence[int]) -> nn.Module: """Create deconvolutional layers by given parameters.""" layers = [] for out_channels, kernel_size in zip(layer_out_channels, layer_kernel_sizes): if kernel_size == 4: padding = 1 output_padding = 0 elif kernel_size == 3: padding = 1 output_padding = 1 elif kernel_size == 2: padding = 0 output_padding = 0 else: raise ValueError(f'Unsupported kernel size {kernel_size} for' 'deconvlutional layers in ' f'{self.__class__.__name__}') cfg = dict( type='deconv', in_channels=in_channels, out_channels=out_channels, kernel_size=kernel_size, stride=2, padding=padding, output_padding=output_padding, bias=False) layers.append(build_upsample_layer(cfg)) layers.append(nn.BatchNorm2d(num_features=out_channels)) layers.append(nn.ReLU(inplace=True)) in_channels = out_channels return nn.Sequential(*layers)
[文档] def forward(self, feats: Sequence[Tensor]) -> List[Tensor]: """Forward the network. The input is multi-stage feature maps and the output is a list of heatmaps from multiple stages. Args: feats (Sequence[Tensor]): Multi-stage feature maps. Returns: List[Tensor]: A list of output heatmaps from multiple stages. """ out = [] assert len(feats) == self.num_stages, ( f'The length of feature maps did not match the ' f'`num_stages` in {self.__class__.__name__}') for i in range(self.num_stages): y = self.multi_deconv_layers[i](feats[i]) y = self.multi_final_layers[i](y) out.append(y) return out
[文档] def predict(self, feats: Features, batch_data_samples: OptSampleList, test_cfg: OptConfigType = {}) -> Predictions: """Predict results from multi-stage feature maps. Args: feats (Tuple[Tensor] | List[Tuple[Tensor]]): The multi-stage features (or multiple multi-stage features in TTA) batch_data_samples (List[:obj:`PoseDataSample`]): The batch data samples test_cfg (dict): The runtime config for testing process. Defaults to {} Returns: Union[InstanceList | Tuple[InstanceList | PixelDataList]]: If ``test_cfg['output_heatmap']==True``, return both pose and heatmap prediction; otherwise only return the pose prediction. The pose prediction is a list of ``InstanceData``, each contains the following fields: - keypoints (np.ndarray): predicted keypoint coordinates in shape (num_instances, K, D) where K is the keypoint number and D is the keypoint dimension - keypoint_scores (np.ndarray): predicted keypoint scores in shape (num_instances, K) The heatmap prediction is a list of ``PixelData``, each contains the following fields: - heatmaps (Tensor): The predicted heatmaps in shape (K, h, w) """ if test_cfg.get('flip_test', False): # TTA: flip test assert isinstance(feats, list) and len(feats) == 2 flip_indices = batch_data_samples[0].metainfo['flip_indices'] _feats, _feats_flip = feats _batch_heatmaps = self.forward(_feats)[-1] _batch_heatmaps_flip = flip_heatmaps( self.forward(_feats_flip)[-1], flip_mode=test_cfg.get('flip_mode', 'heatmap'), flip_indices=flip_indices, shift_heatmap=test_cfg.get('shift_heatmap', False)) batch_heatmaps = (_batch_heatmaps + _batch_heatmaps_flip) * 0.5 else: multi_stage_heatmaps = self.forward(feats) batch_heatmaps = multi_stage_heatmaps[-1] preds = self.decode(batch_heatmaps) if test_cfg.get('output_heatmaps', False): pred_fields = [ PixelData(heatmaps=hm) for hm in batch_heatmaps.detach() ] return preds, pred_fields else: return preds
[文档] def loss(self, feats: Sequence[Tensor], batch_data_samples: OptSampleList, train_cfg: OptConfigType = {}) -> dict: """Calculate losses from a batch of inputs and data samples. Args: feats (Sequence[Tensor]): Multi-stage feature maps. batch_data_samples (List[:obj:`PoseDataSample`]): The Data Samples. It usually includes information such as `gt_instances`. train_cfg (Config, optional): The training config. Returns: dict: A dictionary of loss components. """ multi_stage_pred_heatmaps = self.forward(feats) gt_heatmaps = torch.stack( [d.gt_fields.heatmaps for d in batch_data_samples]) keypoint_weights =[ d.gt_instance_labels.keypoint_weights for d in batch_data_samples ]) # calculate losses over multiple stages losses = dict() for i in range(self.num_stages): if isinstance(self.loss_module, nn.ModuleList): # use different loss_module over different stages loss_func = self.loss_module[i] else: # use the same loss_module over different stages loss_func = self.loss_module # the `gt_heatmaps` and `keypoint_weights` used to calculate loss # for different stages are the same loss_i = loss_func(multi_stage_pred_heatmaps[i], gt_heatmaps, keypoint_weights) if 'loss_kpt' not in losses: losses['loss_kpt'] = loss_i else: losses['loss_kpt'] += loss_i # calculate accuracy _, avg_acc, _ = pose_pck_accuracy( output=to_numpy(multi_stage_pred_heatmaps[-1]), target=to_numpy(gt_heatmaps), mask=to_numpy(keypoint_weights) > 0) acc_pose = torch.tensor(avg_acc, device=gt_heatmaps.device) losses.update(acc_pose=acc_pose) return losses
Read the Docs v: latest
On Read the Docs
Project Home

Free document hosting provided by Read the Docs.