| | import torch |
| | import torch.nn as nn |
| | import torch.nn.functional as F |
| | from mmcv.cnn import ConvModule |
| |
|
| | from mmseg.ops import resize |
| | from ..builder import HEADS |
| | from ..utils import SelfAttentionBlock as _SelfAttentionBlock |
| | from .cascade_decode_head import BaseCascadeDecodeHead |
| |
|
| |
|
| | class SpatialGatherModule(nn.Module): |
| | """Aggregate the context features according to the initial predicted |
| | probability distribution. |
| | |
| | Employ the soft-weighted method to aggregate the context. |
| | """ |
| |
|
| | def __init__(self, scale): |
| | super(SpatialGatherModule, self).__init__() |
| | self.scale = scale |
| |
|
| | def forward(self, feats, probs): |
| | """Forward function.""" |
| | batch_size, num_classes, height, width = probs.size() |
| | channels = feats.size(1) |
| | probs = probs.view(batch_size, num_classes, -1) |
| | feats = feats.view(batch_size, channels, -1) |
| | |
| | feats = feats.permute(0, 2, 1) |
| | |
| | probs = F.softmax(self.scale * probs, dim=2) |
| | |
| | ocr_context = torch.matmul(probs, feats) |
| | ocr_context = ocr_context.permute(0, 2, 1).contiguous().unsqueeze(3) |
| | return ocr_context |
| |
|
| |
|
| | class ObjectAttentionBlock(_SelfAttentionBlock): |
| | """Make a OCR used SelfAttentionBlock.""" |
| |
|
| | def __init__(self, in_channels, channels, scale, conv_cfg, norm_cfg, |
| | act_cfg): |
| | if scale > 1: |
| | query_downsample = nn.MaxPool2d(kernel_size=scale) |
| | else: |
| | query_downsample = None |
| | super(ObjectAttentionBlock, self).__init__( |
| | key_in_channels=in_channels, |
| | query_in_channels=in_channels, |
| | channels=channels, |
| | out_channels=in_channels, |
| | share_key_query=False, |
| | query_downsample=query_downsample, |
| | key_downsample=None, |
| | key_query_num_convs=2, |
| | key_query_norm=True, |
| | value_out_num_convs=1, |
| | value_out_norm=True, |
| | matmul_norm=True, |
| | with_out=True, |
| | conv_cfg=conv_cfg, |
| | norm_cfg=norm_cfg, |
| | act_cfg=act_cfg) |
| | self.bottleneck = ConvModule( |
| | in_channels * 2, |
| | in_channels, |
| | 1, |
| | conv_cfg=self.conv_cfg, |
| | norm_cfg=self.norm_cfg, |
| | act_cfg=self.act_cfg) |
| |
|
| | def forward(self, query_feats, key_feats): |
| | """Forward function.""" |
| | context = super(ObjectAttentionBlock, |
| | self).forward(query_feats, key_feats) |
| | output = self.bottleneck(torch.cat([context, query_feats], dim=1)) |
| | if self.query_downsample is not None: |
| | output = resize(query_feats) |
| |
|
| | return output |
| |
|
| |
|
| | @HEADS.register_module() |
| | class OCRHead(BaseCascadeDecodeHead): |
| | """Object-Contextual Representations for Semantic Segmentation. |
| | |
| | This head is the implementation of `OCRNet |
| | <https://arxiv.org/abs/1909.11065>`_. |
| | |
| | Args: |
| | ocr_channels (int): The intermediate channels of OCR block. |
| | scale (int): The scale of probability map in SpatialGatherModule in |
| | Default: 1. |
| | """ |
| |
|
| | def __init__(self, ocr_channels, scale=1, **kwargs): |
| | super(OCRHead, self).__init__(**kwargs) |
| | self.ocr_channels = ocr_channels |
| | self.scale = scale |
| | self.object_context_block = ObjectAttentionBlock( |
| | self.channels, |
| | self.ocr_channels, |
| | self.scale, |
| | conv_cfg=self.conv_cfg, |
| | norm_cfg=self.norm_cfg, |
| | act_cfg=self.act_cfg) |
| | self.spatial_gather_module = SpatialGatherModule(self.scale) |
| |
|
| | self.bottleneck = ConvModule( |
| | self.in_channels, |
| | self.channels, |
| | 3, |
| | padding=1, |
| | conv_cfg=self.conv_cfg, |
| | norm_cfg=self.norm_cfg, |
| | act_cfg=self.act_cfg) |
| |
|
| | def forward(self, inputs, prev_output): |
| | """Forward function.""" |
| | x = self._transform_inputs(inputs) |
| | feats = self.bottleneck(x) |
| | context = self.spatial_gather_module(feats, prev_output) |
| | object_context = self.object_context_block(feats, context) |
| | output = self.cls_seg(object_context) |
| |
|
| | return output |
| |
|
