| | import torch |
| | import torch.nn as nn |
| | import torch.nn.functional as F |
| | from annotator.mmpkg.mmcv.cnn import ConvModule, build_norm_layer |
| |
|
| | from annotator.mmpkg.mmseg.ops import Encoding, resize |
| | from ..builder import HEADS, build_loss |
| | from .decode_head import BaseDecodeHead |
| |
|
| |
|
| | class EncModule(nn.Module): |
| | """Encoding Module used in EncNet. |
| | |
| | Args: |
| | in_channels (int): Input channels. |
| | num_codes (int): Number of code words. |
| | conv_cfg (dict|None): Config of conv layers. |
| | norm_cfg (dict|None): Config of norm layers. |
| | act_cfg (dict): Config of activation layers. |
| | """ |
| |
|
| | def __init__(self, in_channels, num_codes, conv_cfg, norm_cfg, act_cfg): |
| | super(EncModule, self).__init__() |
| | self.encoding_project = ConvModule( |
| | in_channels, |
| | in_channels, |
| | 1, |
| | conv_cfg=conv_cfg, |
| | norm_cfg=norm_cfg, |
| | act_cfg=act_cfg) |
| | |
| | |
| | if norm_cfg is not None: |
| | encoding_norm_cfg = norm_cfg.copy() |
| | if encoding_norm_cfg['type'] in ['BN', 'IN']: |
| | encoding_norm_cfg['type'] += '1d' |
| | else: |
| | encoding_norm_cfg['type'] = encoding_norm_cfg['type'].replace( |
| | '2d', '1d') |
| | else: |
| | |
| | encoding_norm_cfg = dict(type='BN1d') |
| | self.encoding = nn.Sequential( |
| | Encoding(channels=in_channels, num_codes=num_codes), |
| | build_norm_layer(encoding_norm_cfg, num_codes)[1], |
| | nn.ReLU(inplace=True)) |
| | self.fc = nn.Sequential( |
| | nn.Linear(in_channels, in_channels), nn.Sigmoid()) |
| |
|
| | def forward(self, x): |
| | """Forward function.""" |
| | encoding_projection = self.encoding_project(x) |
| | encoding_feat = self.encoding(encoding_projection).mean(dim=1) |
| | batch_size, channels, _, _ = x.size() |
| | gamma = self.fc(encoding_feat) |
| | y = gamma.view(batch_size, channels, 1, 1) |
| | output = F.relu_(x + x * y) |
| | return encoding_feat, output |
| |
|
| |
|
| | @HEADS.register_module() |
| | class EncHead(BaseDecodeHead): |
| | """Context Encoding for Semantic Segmentation. |
| | |
| | This head is the implementation of `EncNet |
| | <https://arxiv.org/abs/1803.08904>`_. |
| | |
| | Args: |
| | num_codes (int): Number of code words. Default: 32. |
| | use_se_loss (bool): Whether use Semantic Encoding Loss (SE-loss) to |
| | regularize the training. Default: True. |
| | add_lateral (bool): Whether use lateral connection to fuse features. |
| | Default: False. |
| | loss_se_decode (dict): Config of decode loss. |
| | Default: dict(type='CrossEntropyLoss', use_sigmoid=True). |
| | """ |
| |
|
| | def __init__(self, |
| | num_codes=32, |
| | use_se_loss=True, |
| | add_lateral=False, |
| | loss_se_decode=dict( |
| | type='CrossEntropyLoss', |
| | use_sigmoid=True, |
| | loss_weight=0.2), |
| | **kwargs): |
| | super(EncHead, self).__init__( |
| | input_transform='multiple_select', **kwargs) |
| | self.use_se_loss = use_se_loss |
| | self.add_lateral = add_lateral |
| | self.num_codes = num_codes |
| | self.bottleneck = ConvModule( |
| | self.in_channels[-1], |
| | self.channels, |
| | 3, |
| | padding=1, |
| | conv_cfg=self.conv_cfg, |
| | norm_cfg=self.norm_cfg, |
| | act_cfg=self.act_cfg) |
| | if add_lateral: |
| | self.lateral_convs = nn.ModuleList() |
| | for in_channels in self.in_channels[:-1]: |
| | self.lateral_convs.append( |
| | ConvModule( |
| | in_channels, |
| | self.channels, |
| | 1, |
| | conv_cfg=self.conv_cfg, |
| | norm_cfg=self.norm_cfg, |
| | act_cfg=self.act_cfg)) |
| | self.fusion = ConvModule( |
| | len(self.in_channels) * self.channels, |
| | self.channels, |
| | 3, |
| | padding=1, |
| | conv_cfg=self.conv_cfg, |
| | norm_cfg=self.norm_cfg, |
| | act_cfg=self.act_cfg) |
| | self.enc_module = EncModule( |
| | self.channels, |
| | num_codes=num_codes, |
| | conv_cfg=self.conv_cfg, |
| | norm_cfg=self.norm_cfg, |
| | act_cfg=self.act_cfg) |
| | if self.use_se_loss: |
| | self.loss_se_decode = build_loss(loss_se_decode) |
| | self.se_layer = nn.Linear(self.channels, self.num_classes) |
| |
|
| | def forward(self, inputs): |
| | """Forward function.""" |
| | inputs = self._transform_inputs(inputs) |
| | feat = self.bottleneck(inputs[-1]) |
| | if self.add_lateral: |
| | laterals = [ |
| | resize( |
| | lateral_conv(inputs[i]), |
| | size=feat.shape[2:], |
| | mode='bilinear', |
| | align_corners=self.align_corners) |
| | for i, lateral_conv in enumerate(self.lateral_convs) |
| | ] |
| | feat = self.fusion(torch.cat([feat, *laterals], 1)) |
| | encode_feat, output = self.enc_module(feat) |
| | output = self.cls_seg(output) |
| | if self.use_se_loss: |
| | se_output = self.se_layer(encode_feat) |
| | return output, se_output |
| | else: |
| | return output |
| |
|
| | def forward_test(self, inputs, img_metas, test_cfg): |
| | """Forward function for testing, ignore se_loss.""" |
| | if self.use_se_loss: |
| | return self.forward(inputs)[0] |
| | else: |
| | return self.forward(inputs) |
| |
|
| | @staticmethod |
| | def _convert_to_onehot_labels(seg_label, num_classes): |
| | """Convert segmentation label to onehot. |
| | |
| | Args: |
| | seg_label (Tensor): Segmentation label of shape (N, H, W). |
| | num_classes (int): Number of classes. |
| | |
| | Returns: |
| | Tensor: Onehot labels of shape (N, num_classes). |
| | """ |
| |
|
| | batch_size = seg_label.size(0) |
| | onehot_labels = seg_label.new_zeros((batch_size, num_classes)) |
| | for i in range(batch_size): |
| | hist = seg_label[i].float().histc( |
| | bins=num_classes, min=0, max=num_classes - 1) |
| | onehot_labels[i] = hist > 0 |
| | return onehot_labels |
| |
|
| | def losses(self, seg_logit, seg_label): |
| | """Compute segmentation and semantic encoding loss.""" |
| | seg_logit, se_seg_logit = seg_logit |
| | loss = dict() |
| | loss.update(super(EncHead, self).losses(seg_logit, seg_label)) |
| | se_loss = self.loss_se_decode( |
| | se_seg_logit, |
| | self._convert_to_onehot_labels(seg_label, self.num_classes)) |
| | loss['loss_se'] = se_loss |
| | return loss |
| |
|
