Source code for mmpose.models.necks.fmap_proc_neck

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

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch import Tensor

from mmpose.models.utils.ops import resize
from mmpose.registry import MODELS


@MODELS.register_module()
class FeatureMapProcessor(nn.Module):
    """A PyTorch module for selecting, concatenating, and rescaling feature
    maps.

    Args:
        select_index (Optional[Union[int, Tuple[int]]], optional): Index or
            indices of feature maps to select. Defaults to None, which means
            all feature maps are used.
        concat (bool, optional): Whether to concatenate the selected feature
            maps. Defaults to False.
        scale_factor (float, optional): The scaling factor to apply to the
            feature maps. Defaults to 1.0.
        apply_relu (bool, optional): Whether to apply ReLU on input feature
            maps. Defaults to False.
        align_corners (bool, optional): Whether to align corners when resizing
            the feature maps. Defaults to False.
    """

    def __init__(
        self,
        select_index: Optional[Union[int, Tuple[int]]] = None,
        concat: bool = False,
        scale_factor: float = 1.0,
        apply_relu: bool = False,
        align_corners: bool = False,
    ):
        super().__init__()

        if isinstance(select_index, int):
            select_index = (select_index, )
        self.select_index = select_index
        self.concat = concat

        assert (
            scale_factor > 0
        ), f'the argument `scale_factor` must be positive, ' \
           f'but got {scale_factor}'
        self.scale_factor = scale_factor
        self.apply_relu = apply_relu
        self.align_corners = align_corners

    def forward(self, inputs: Union[Tensor, Sequence[Tensor]]
                ) -> Union[Tensor, List[Tensor]]:

        if not isinstance(inputs, (tuple, list)):
            sequential_input = False
            inputs = [inputs]
        else:
            sequential_input = True

            if self.select_index is not None:
                inputs = [inputs[i] for i in self.select_index]

            if self.concat:
                inputs = self._concat(inputs)

        if self.apply_relu:
            inputs = [F.relu(x) for x in inputs]

        if self.scale_factor != 1.0:
            inputs = self._rescale(inputs)

        if not sequential_input:
            inputs = inputs[0]

        return inputs

    def _concat(self, inputs: Sequence[Tensor]) -> List[Tensor]:
        size = inputs[0].shape[-2:]
        resized_inputs = [
            resize(
                x,
                size=size,
                mode='bilinear',
                align_corners=self.align_corners) for x in inputs
        ]
        return [torch.cat(resized_inputs, dim=1)]

    def _rescale(self, inputs: Sequence[Tensor]) -> List[Tensor]:
        rescaled_inputs = [
            resize(
                x,
                scale_factor=self.scale_factor,
                mode='bilinear',
                align_corners=self.align_corners,
            ) for x in inputs
        ]
        return rescaled_inputs