| | |
| | from typing import Tuple |
| |
|
| | import torch.nn.functional as F |
| | from mmcv.cnn import ConvModule |
| | from mmcv.cnn.bricks import NonLocal2d |
| | from mmengine.model import BaseModule |
| | from torch import Tensor |
| |
|
| | from mmdet.registry import MODELS |
| | from mmdet.utils import OptConfigType, OptMultiConfig |
| |
|
| |
|
| | @MODELS.register_module() |
| | class BFP(BaseModule): |
| | """BFP (Balanced Feature Pyramids) |
| | |
| | BFP takes multi-level features as inputs and gather them into a single one, |
| | then refine the gathered feature and scatter the refined results to |
| | multi-level features. This module is used in Libra R-CNN (CVPR 2019), see |
| | the paper `Libra R-CNN: Towards Balanced Learning for Object Detection |
| | <https://arxiv.org/abs/1904.02701>`_ for details. |
| | |
| | Args: |
| | in_channels (int): Number of input channels (feature maps of all levels |
| | should have the same channels). |
| | num_levels (int): Number of input feature levels. |
| | refine_level (int): Index of integration and refine level of BSF in |
| | multi-level features from bottom to top. |
| | refine_type (str): Type of the refine op, currently support |
| | [None, 'conv', 'non_local']. |
| | conv_cfg (:obj:`ConfigDict` or dict, optional): The config dict for |
| | convolution layers. |
| | norm_cfg (:obj:`ConfigDict` or dict, optional): The config dict for |
| | normalization layers. |
| | init_cfg (:obj:`ConfigDict` or dict or list[:obj:`ConfigDict` or |
| | dict], optional): Initialization config dict. |
| | """ |
| |
|
| | def __init__( |
| | self, |
| | in_channels: int, |
| | num_levels: int, |
| | refine_level: int = 2, |
| | refine_type: str = None, |
| | conv_cfg: OptConfigType = None, |
| | norm_cfg: OptConfigType = None, |
| | init_cfg: OptMultiConfig = dict( |
| | type='Xavier', layer='Conv2d', distribution='uniform') |
| | ) -> None: |
| | super().__init__(init_cfg=init_cfg) |
| | assert refine_type in [None, 'conv', 'non_local'] |
| |
|
| | self.in_channels = in_channels |
| | self.num_levels = num_levels |
| | self.conv_cfg = conv_cfg |
| | self.norm_cfg = norm_cfg |
| |
|
| | self.refine_level = refine_level |
| | self.refine_type = refine_type |
| | assert 0 <= self.refine_level < self.num_levels |
| |
|
| | if self.refine_type == 'conv': |
| | self.refine = ConvModule( |
| | self.in_channels, |
| | self.in_channels, |
| | 3, |
| | padding=1, |
| | conv_cfg=self.conv_cfg, |
| | norm_cfg=self.norm_cfg) |
| | elif self.refine_type == 'non_local': |
| | self.refine = NonLocal2d( |
| | self.in_channels, |
| | reduction=1, |
| | use_scale=False, |
| | conv_cfg=self.conv_cfg, |
| | norm_cfg=self.norm_cfg) |
| |
|
| | def forward(self, inputs: Tuple[Tensor]) -> Tuple[Tensor]: |
| | """Forward function.""" |
| | assert len(inputs) == self.num_levels |
| |
|
| | |
| | feats = [] |
| | gather_size = inputs[self.refine_level].size()[2:] |
| | for i in range(self.num_levels): |
| | if i < self.refine_level: |
| | gathered = F.adaptive_max_pool2d( |
| | inputs[i], output_size=gather_size) |
| | else: |
| | gathered = F.interpolate( |
| | inputs[i], size=gather_size, mode='nearest') |
| | feats.append(gathered) |
| |
|
| | bsf = sum(feats) / len(feats) |
| |
|
| | |
| | if self.refine_type is not None: |
| | bsf = self.refine(bsf) |
| |
|
| | |
| | outs = [] |
| | for i in range(self.num_levels): |
| | out_size = inputs[i].size()[2:] |
| | if i < self.refine_level: |
| | residual = F.interpolate(bsf, size=out_size, mode='nearest') |
| | else: |
| | residual = F.adaptive_max_pool2d(bsf, output_size=out_size) |
| | outs.append(residual + inputs[i]) |
| |
|
| | return tuple(outs) |
| |
|
