|
import numpy as np |
|
import torch.nn as nn |
|
from mmcv.cnn import ConvModule |
|
|
|
from mmseg.ops import resize |
|
from ..builder import HEADS |
|
from .decode_head import BaseDecodeHead |
|
|
|
|
|
@HEADS.register_module() |
|
class FPNHead(BaseDecodeHead): |
|
"""Panoptic Feature Pyramid Networks. |
|
|
|
This head is the implementation of `Semantic FPN |
|
<https://arxiv.org/abs/1901.02446>`_. |
|
|
|
Args: |
|
feature_strides (tuple[int]): The strides for input feature maps. |
|
stack_lateral. All strides suppose to be power of 2. The first |
|
one is of largest resolution. |
|
""" |
|
|
|
def __init__(self, feature_strides, **kwargs): |
|
super(FPNHead, self).__init__( |
|
input_transform='multiple_select', **kwargs) |
|
assert len(feature_strides) == len(self.in_channels) |
|
assert min(feature_strides) == feature_strides[0] |
|
self.feature_strides = feature_strides |
|
|
|
self.scale_heads = nn.ModuleList() |
|
for i in range(len(feature_strides)): |
|
head_length = max( |
|
1, |
|
int(np.log2(feature_strides[i]) - np.log2(feature_strides[0]))) |
|
scale_head = [] |
|
for k in range(head_length): |
|
scale_head.append( |
|
ConvModule( |
|
self.in_channels[i] if k == 0 else self.channels, |
|
self.channels, |
|
3, |
|
padding=1, |
|
conv_cfg=self.conv_cfg, |
|
norm_cfg=self.norm_cfg, |
|
act_cfg=self.act_cfg)) |
|
if feature_strides[i] != feature_strides[0]: |
|
scale_head.append( |
|
nn.Upsample( |
|
scale_factor=2, |
|
mode='bilinear', |
|
align_corners=self.align_corners)) |
|
self.scale_heads.append(nn.Sequential(*scale_head)) |
|
|
|
def forward(self, inputs): |
|
|
|
x = self._transform_inputs(inputs) |
|
|
|
output = self.scale_heads[0](x[0]) |
|
for i in range(1, len(self.feature_strides)): |
|
|
|
output = output + resize( |
|
self.scale_heads[i](x[i]), |
|
size=output.shape[2:], |
|
mode='bilinear', |
|
align_corners=self.align_corners) |
|
|
|
output = self.cls_seg(output) |
|
return output |
|
|
