| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
| from annotator.mmpkg.mmcv.cnn import ConvModule |
|
|
| from annotator.mmpkg.mmseg.ops import resize |
| from ..builder import HEADS |
| from .decode_head import BaseDecodeHead |
|
|
| try: |
| try: |
| from mmcv.ops import PSAMask |
| except ImportError: |
| from annotator.mmpkg.mmcv.ops import PSAMask |
| except ModuleNotFoundError: |
| PSAMask = None |
|
|
|
|
| @HEADS.register_module() |
| class PSAHead(BaseDecodeHead): |
| """Point-wise Spatial Attention Network for Scene Parsing. |
| |
| This head is the implementation of `PSANet |
| <https://hszhao.github.io/papers/eccv18_psanet.pdf>`_. |
| |
| Args: |
| mask_size (tuple[int]): The PSA mask size. It usually equals input |
| size. |
| psa_type (str): The type of psa module. Options are 'collect', |
| 'distribute', 'bi-direction'. Default: 'bi-direction' |
| compact (bool): Whether use compact map for 'collect' mode. |
| Default: True. |
| shrink_factor (int): The downsample factors of psa mask. Default: 2. |
| normalization_factor (float): The normalize factor of attention. |
| psa_softmax (bool): Whether use softmax for attention. |
| """ |
|
|
| def __init__(self, |
| mask_size, |
| psa_type='bi-direction', |
| compact=False, |
| shrink_factor=2, |
| normalization_factor=1.0, |
| psa_softmax=True, |
| **kwargs): |
| if PSAMask is None: |
| raise RuntimeError('Please install mmcv-full for PSAMask ops') |
| super(PSAHead, self).__init__(**kwargs) |
| assert psa_type in ['collect', 'distribute', 'bi-direction'] |
| self.psa_type = psa_type |
| self.compact = compact |
| self.shrink_factor = shrink_factor |
| self.mask_size = mask_size |
| mask_h, mask_w = mask_size |
| self.psa_softmax = psa_softmax |
| if normalization_factor is None: |
| normalization_factor = mask_h * mask_w |
| self.normalization_factor = normalization_factor |
|
|
| self.reduce = ConvModule( |
| self.in_channels, |
| self.channels, |
| kernel_size=1, |
| conv_cfg=self.conv_cfg, |
| norm_cfg=self.norm_cfg, |
| act_cfg=self.act_cfg) |
| self.attention = nn.Sequential( |
| ConvModule( |
| self.channels, |
| self.channels, |
| kernel_size=1, |
| conv_cfg=self.conv_cfg, |
| norm_cfg=self.norm_cfg, |
| act_cfg=self.act_cfg), |
| nn.Conv2d( |
| self.channels, mask_h * mask_w, kernel_size=1, bias=False)) |
| if psa_type == 'bi-direction': |
| self.reduce_p = ConvModule( |
| self.in_channels, |
| self.channels, |
| kernel_size=1, |
| conv_cfg=self.conv_cfg, |
| norm_cfg=self.norm_cfg, |
| act_cfg=self.act_cfg) |
| self.attention_p = nn.Sequential( |
| ConvModule( |
| self.channels, |
| self.channels, |
| kernel_size=1, |
| conv_cfg=self.conv_cfg, |
| norm_cfg=self.norm_cfg, |
| act_cfg=self.act_cfg), |
| nn.Conv2d( |
| self.channels, mask_h * mask_w, kernel_size=1, bias=False)) |
| self.psamask_collect = PSAMask('collect', mask_size) |
| self.psamask_distribute = PSAMask('distribute', mask_size) |
| else: |
| self.psamask = PSAMask(psa_type, mask_size) |
| self.proj = ConvModule( |
| self.channels * (2 if psa_type == 'bi-direction' else 1), |
| self.in_channels, |
| kernel_size=1, |
| padding=1, |
| conv_cfg=self.conv_cfg, |
| norm_cfg=self.norm_cfg, |
| act_cfg=self.act_cfg) |
| self.bottleneck = ConvModule( |
| self.in_channels * 2, |
| self.channels, |
| kernel_size=3, |
| padding=1, |
| conv_cfg=self.conv_cfg, |
| norm_cfg=self.norm_cfg, |
| act_cfg=self.act_cfg) |
|
|
| def forward(self, inputs): |
| """Forward function.""" |
| x = self._transform_inputs(inputs) |
| identity = x |
| align_corners = self.align_corners |
| if self.psa_type in ['collect', 'distribute']: |
| out = self.reduce(x) |
| n, c, h, w = out.size() |
| if self.shrink_factor != 1: |
| if h % self.shrink_factor and w % self.shrink_factor: |
| h = (h - 1) // self.shrink_factor + 1 |
| w = (w - 1) // self.shrink_factor + 1 |
| align_corners = True |
| else: |
| h = h // self.shrink_factor |
| w = w // self.shrink_factor |
| align_corners = False |
| out = resize( |
| out, |
| size=(h, w), |
| mode='bilinear', |
| align_corners=align_corners) |
| y = self.attention(out) |
| if self.compact: |
| if self.psa_type == 'collect': |
| y = y.view(n, h * w, |
| h * w).transpose(1, 2).view(n, h * w, h, w) |
| else: |
| y = self.psamask(y) |
| if self.psa_softmax: |
| y = F.softmax(y, dim=1) |
| out = torch.bmm( |
| out.view(n, c, h * w), y.view(n, h * w, h * w)).view( |
| n, c, h, w) * (1.0 / self.normalization_factor) |
| else: |
| x_col = self.reduce(x) |
| x_dis = self.reduce_p(x) |
| n, c, h, w = x_col.size() |
| if self.shrink_factor != 1: |
| if h % self.shrink_factor and w % self.shrink_factor: |
| h = (h - 1) // self.shrink_factor + 1 |
| w = (w - 1) // self.shrink_factor + 1 |
| align_corners = True |
| else: |
| h = h // self.shrink_factor |
| w = w // self.shrink_factor |
| align_corners = False |
| x_col = resize( |
| x_col, |
| size=(h, w), |
| mode='bilinear', |
| align_corners=align_corners) |
| x_dis = resize( |
| x_dis, |
| size=(h, w), |
| mode='bilinear', |
| align_corners=align_corners) |
| y_col = self.attention(x_col) |
| y_dis = self.attention_p(x_dis) |
| if self.compact: |
| y_dis = y_dis.view(n, h * w, |
| h * w).transpose(1, 2).view(n, h * w, h, w) |
| else: |
| y_col = self.psamask_collect(y_col) |
| y_dis = self.psamask_distribute(y_dis) |
| if self.psa_softmax: |
| y_col = F.softmax(y_col, dim=1) |
| y_dis = F.softmax(y_dis, dim=1) |
| x_col = torch.bmm( |
| x_col.view(n, c, h * w), y_col.view(n, h * w, h * w)).view( |
| n, c, h, w) * (1.0 / self.normalization_factor) |
| x_dis = torch.bmm( |
| x_dis.view(n, c, h * w), y_dis.view(n, h * w, h * w)).view( |
| n, c, h, w) * (1.0 / self.normalization_factor) |
| out = torch.cat([x_col, x_dis], 1) |
| out = self.proj(out) |
| out = resize( |
| out, |
| size=identity.shape[2:], |
| mode='bilinear', |
| align_corners=align_corners) |
| out = self.bottleneck(torch.cat((identity, out), dim=1)) |
| out = self.cls_seg(out) |
| return out |
|
|
