|
import torch |
|
import torch.nn as nn |
|
from annotator.uniformer.mmcv.cnn import ConvModule |
|
|
|
from ..builder import HEADS |
|
from .decode_head import BaseDecodeHead |
|
|
|
|
|
@HEADS.register_module() |
|
class FCNHead(BaseDecodeHead): |
|
"""Fully Convolution Networks for Semantic Segmentation. |
|
|
|
This head is implemented of `FCNNet <https://arxiv.org/abs/1411.4038>`_. |
|
|
|
Args: |
|
num_convs (int): Number of convs in the head. Default: 2. |
|
kernel_size (int): The kernel size for convs in the head. Default: 3. |
|
concat_input (bool): Whether concat the input and output of convs |
|
before classification layer. |
|
dilation (int): The dilation rate for convs in the head. Default: 1. |
|
""" |
|
|
|
def __init__(self, |
|
num_convs=2, |
|
kernel_size=3, |
|
concat_input=True, |
|
dilation=1, |
|
**kwargs): |
|
assert num_convs >= 0 and dilation > 0 and isinstance(dilation, int) |
|
self.num_convs = num_convs |
|
self.concat_input = concat_input |
|
self.kernel_size = kernel_size |
|
super(FCNHead, self).__init__(**kwargs) |
|
if num_convs == 0: |
|
assert self.in_channels == self.channels |
|
|
|
conv_padding = (kernel_size // 2) * dilation |
|
convs = [] |
|
convs.append( |
|
ConvModule( |
|
self.in_channels, |
|
self.channels, |
|
kernel_size=kernel_size, |
|
padding=conv_padding, |
|
dilation=dilation, |
|
conv_cfg=self.conv_cfg, |
|
norm_cfg=self.norm_cfg, |
|
act_cfg=self.act_cfg)) |
|
for i in range(num_convs - 1): |
|
convs.append( |
|
ConvModule( |
|
self.channels, |
|
self.channels, |
|
kernel_size=kernel_size, |
|
padding=conv_padding, |
|
dilation=dilation, |
|
conv_cfg=self.conv_cfg, |
|
norm_cfg=self.norm_cfg, |
|
act_cfg=self.act_cfg)) |
|
if num_convs == 0: |
|
self.convs = nn.Identity() |
|
else: |
|
self.convs = nn.Sequential(*convs) |
|
if self.concat_input: |
|
self.conv_cat = ConvModule( |
|
self.in_channels + self.channels, |
|
self.channels, |
|
kernel_size=kernel_size, |
|
padding=kernel_size // 2, |
|
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) |
|
output = self.convs(x) |
|
if self.concat_input: |
|
output = self.conv_cat(torch.cat([x, output], dim=1)) |
|
output = self.cls_seg(output) |
|
return output |
|
|
