|
from abc import ABCMeta, abstractmethod |
|
|
|
import torch |
|
import torch.nn as nn |
|
from mmcv.cnn import normal_init |
|
from mmcv.runner import auto_fp16, force_fp32 |
|
|
|
from mmseg.core import build_pixel_sampler |
|
from mmseg.ops import resize |
|
from ..builder import build_loss |
|
from ..losses import accuracy |
|
|
|
|
|
class BaseDecodeHead(nn.Module, metaclass=ABCMeta): |
|
"""Base class for BaseDecodeHead. |
|
|
|
Args: |
|
in_channels (int|Sequence[int]): Input channels. |
|
channels (int): Channels after modules, before conv_seg. |
|
num_classes (int): Number of classes. |
|
dropout_ratio (float): Ratio of dropout layer. Default: 0.1. |
|
conv_cfg (dict|None): Config of conv layers. Default: None. |
|
norm_cfg (dict|None): Config of norm layers. Default: None. |
|
act_cfg (dict): Config of activation layers. |
|
Default: dict(type='ReLU') |
|
in_index (int|Sequence[int]): Input feature index. Default: -1 |
|
input_transform (str|None): Transformation type of input features. |
|
Options: 'resize_concat', 'multiple_select', None. |
|
'resize_concat': Multiple feature maps will be resize to the |
|
same size as first one and than concat together. |
|
Usually used in FCN head of HRNet. |
|
'multiple_select': Multiple feature maps will be bundle into |
|
a list and passed into decode head. |
|
None: Only one select feature map is allowed. |
|
Default: None. |
|
loss_decode (dict): Config of decode loss. |
|
Default: dict(type='CrossEntropyLoss'). |
|
ignore_index (int | None): The label index to be ignored. When using |
|
masked BCE loss, ignore_index should be set to None. Default: 255 |
|
sampler (dict|None): The config of segmentation map sampler. |
|
Default: None. |
|
align_corners (bool): align_corners argument of F.interpolate. |
|
Default: False. |
|
""" |
|
|
|
def __init__(self, |
|
in_channels, |
|
channels, |
|
*, |
|
num_classes, |
|
dropout_ratio=0.1, |
|
conv_cfg=None, |
|
norm_cfg=None, |
|
act_cfg=dict(type='ReLU'), |
|
in_index=-1, |
|
input_transform=None, |
|
loss_decode=dict( |
|
type='CrossEntropyLoss', |
|
use_sigmoid=False, |
|
loss_weight=1.0), |
|
ignore_index=255, |
|
sampler=None, |
|
align_corners=False): |
|
super(BaseDecodeHead, self).__init__() |
|
self._init_inputs(in_channels, in_index, input_transform) |
|
self.channels = channels |
|
self.num_classes = num_classes |
|
self.dropout_ratio = dropout_ratio |
|
self.conv_cfg = conv_cfg |
|
self.norm_cfg = norm_cfg |
|
self.act_cfg = act_cfg |
|
self.in_index = in_index |
|
self.loss_decode = build_loss(loss_decode) |
|
self.ignore_index = ignore_index |
|
self.align_corners = align_corners |
|
if sampler is not None: |
|
self.sampler = build_pixel_sampler(sampler, context=self) |
|
else: |
|
self.sampler = None |
|
|
|
self.conv_seg = nn.Conv2d(channels, num_classes, kernel_size=1) |
|
if dropout_ratio > 0: |
|
self.dropout = nn.Dropout2d(dropout_ratio) |
|
else: |
|
self.dropout = None |
|
self.fp16_enabled = False |
|
|
|
def extra_repr(self): |
|
"""Extra repr.""" |
|
s = f'input_transform={self.input_transform}, ' \ |
|
f'ignore_index={self.ignore_index}, ' \ |
|
f'align_corners={self.align_corners}' |
|
return s |
|
|
|
def _init_inputs(self, in_channels, in_index, input_transform): |
|
"""Check and initialize input transforms. |
|
|
|
The in_channels, in_index and input_transform must match. |
|
Specifically, when input_transform is None, only single feature map |
|
will be selected. So in_channels and in_index must be of type int. |
|
When input_transform |
|
|
|
Args: |
|
in_channels (int|Sequence[int]): Input channels. |
|
in_index (int|Sequence[int]): Input feature index. |
|
input_transform (str|None): Transformation type of input features. |
|
Options: 'resize_concat', 'multiple_select', None. |
|
'resize_concat': Multiple feature maps will be resize to the |
|
same size as first one and than concat together. |
|
Usually used in FCN head of HRNet. |
|
'multiple_select': Multiple feature maps will be bundle into |
|
a list and passed into decode head. |
|
None: Only one select feature map is allowed. |
|
""" |
|
|
|
if input_transform is not None: |
|
assert input_transform in ['resize_concat', 'multiple_select'] |
|
self.input_transform = input_transform |
|
self.in_index = in_index |
|
if input_transform is not None: |
|
assert isinstance(in_channels, (list, tuple)) |
|
assert isinstance(in_index, (list, tuple)) |
|
assert len(in_channels) == len(in_index) |
|
if input_transform == 'resize_concat': |
|
self.in_channels = sum(in_channels) |
|
else: |
|
self.in_channels = in_channels |
|
else: |
|
assert isinstance(in_channels, int) |
|
assert isinstance(in_index, int) |
|
self.in_channels = in_channels |
|
|
|
def init_weights(self): |
|
"""Initialize weights of classification layer.""" |
|
normal_init(self.conv_seg, mean=0, std=0.01) |
|
|
|
def _transform_inputs(self, inputs): |
|
"""Transform inputs for decoder. |
|
|
|
Args: |
|
inputs (list[Tensor]): List of multi-level img features. |
|
|
|
Returns: |
|
Tensor: The transformed inputs |
|
""" |
|
|
|
if self.input_transform == 'resize_concat': |
|
inputs = [inputs[i] for i in self.in_index] |
|
upsampled_inputs = [ |
|
resize( |
|
input=x, |
|
size=inputs[0].shape[2:], |
|
mode='bilinear', |
|
align_corners=self.align_corners) for x in inputs |
|
] |
|
inputs = torch.cat(upsampled_inputs, dim=1) |
|
elif self.input_transform == 'multiple_select': |
|
inputs = [inputs[i] for i in self.in_index] |
|
else: |
|
inputs = inputs[self.in_index] |
|
|
|
return inputs |
|
|
|
@auto_fp16() |
|
@abstractmethod |
|
def forward(self, inputs): |
|
"""Placeholder of forward function.""" |
|
pass |
|
|
|
def forward_train(self, inputs, img_metas, gt_semantic_seg, train_cfg): |
|
"""Forward function for training. |
|
Args: |
|
inputs (list[Tensor]): List of multi-level img features. |
|
img_metas (list[dict]): List of image info dict where each dict |
|
has: 'img_shape', 'scale_factor', 'flip', and may also contain |
|
'filename', 'ori_shape', 'pad_shape', and 'img_norm_cfg'. |
|
For details on the values of these keys see |
|
`mmseg/datasets/pipelines/formatting.py:Collect`. |
|
gt_semantic_seg (Tensor): Semantic segmentation masks |
|
used if the architecture supports semantic segmentation task. |
|
train_cfg (dict): The training config. |
|
|
|
Returns: |
|
dict[str, Tensor]: a dictionary of loss components |
|
""" |
|
seg_logits = self.forward(inputs) |
|
losses = self.losses(seg_logits, gt_semantic_seg) |
|
return losses |
|
|
|
def forward_test(self, inputs, img_metas, test_cfg): |
|
"""Forward function for testing. |
|
|
|
Args: |
|
inputs (list[Tensor]): List of multi-level img features. |
|
img_metas (list[dict]): List of image info dict where each dict |
|
has: 'img_shape', 'scale_factor', 'flip', and may also contain |
|
'filename', 'ori_shape', 'pad_shape', and 'img_norm_cfg'. |
|
For details on the values of these keys see |
|
`mmseg/datasets/pipelines/formatting.py:Collect`. |
|
test_cfg (dict): The testing config. |
|
|
|
Returns: |
|
Tensor: Output segmentation map. |
|
""" |
|
return self.forward(inputs) |
|
|
|
def cls_seg(self, feat): |
|
"""Classify each pixel.""" |
|
if self.dropout is not None: |
|
feat = self.dropout(feat) |
|
output = self.conv_seg(feat) |
|
return output |
|
|
|
@force_fp32(apply_to=('seg_logit', )) |
|
def losses(self, seg_logit, seg_label): |
|
"""Compute segmentation loss.""" |
|
loss = dict() |
|
seg_logit = resize( |
|
input=seg_logit, |
|
size=seg_label.shape[2:], |
|
mode='bilinear', |
|
align_corners=self.align_corners) |
|
if self.sampler is not None: |
|
seg_weight = self.sampler.sample(seg_logit, seg_label) |
|
else: |
|
seg_weight = None |
|
seg_label = seg_label.squeeze(1) |
|
loss['loss_seg'] = self.loss_decode( |
|
seg_logit, |
|
seg_label, |
|
weight=seg_weight, |
|
ignore_index=self.ignore_index) |
|
loss['acc_seg'] = accuracy(seg_logit, seg_label) |
|
return loss |
|
|
