|
|
|
|
|
|
|
|
import torch |
|
|
import torch.nn as nn |
|
|
import torch.nn.functional as F |
|
|
from mmengine.model.weight_init import caffe2_xavier_init, kaiming_init |
|
|
|
|
|
from mmdet.registry import MODELS |
|
|
|
|
|
|
|
|
class PyramidPoolingModule(nn.Module): |
|
|
|
|
|
def __init__(self, |
|
|
in_channels, |
|
|
channels=512, |
|
|
sizes=(1, 2, 3, 6), |
|
|
act_cfg=dict(type='ReLU')): |
|
|
super().__init__() |
|
|
self.stages = [] |
|
|
self.stages = nn.ModuleList( |
|
|
[self._make_stage(in_channels, channels, size) for size in sizes]) |
|
|
self.bottleneck = nn.Conv2d(in_channels + len(sizes) * channels, |
|
|
in_channels, 1) |
|
|
self.act = MODELS.build(act_cfg) |
|
|
|
|
|
def _make_stage(self, features, out_features, size): |
|
|
prior = nn.AdaptiveAvgPool2d(output_size=(size, size)) |
|
|
conv = nn.Conv2d(features, out_features, 1) |
|
|
return nn.Sequential(prior, conv) |
|
|
|
|
|
def forward(self, feats): |
|
|
h, w = feats.size(2), feats.size(3) |
|
|
priors = [ |
|
|
F.interpolate( |
|
|
input=self.act(stage(feats)), |
|
|
size=(h, w), |
|
|
mode='bilinear', |
|
|
align_corners=False) for stage in self.stages |
|
|
] + [feats] |
|
|
out = self.act(self.bottleneck(torch.cat(priors, 1))) |
|
|
return out |
|
|
|
|
|
|
|
|
@MODELS.register_module() |
|
|
class InstanceContextEncoder(nn.Module): |
|
|
""" |
|
|
Instance Context Encoder |
|
|
1. construct feature pyramids from ResNet |
|
|
2. enlarge receptive fields (ppm) |
|
|
3. multi-scale fusion |
|
|
""" |
|
|
|
|
|
def __init__(self, |
|
|
in_channels, |
|
|
out_channels=256, |
|
|
with_ppm=True, |
|
|
act_cfg=dict(type='ReLU')): |
|
|
super().__init__() |
|
|
self.num_channels = out_channels |
|
|
self.in_channels = in_channels |
|
|
self.with_ppm = with_ppm |
|
|
fpn_laterals = [] |
|
|
fpn_outputs = [] |
|
|
for in_channel in reversed(self.in_channels): |
|
|
lateral_conv = nn.Conv2d(in_channel, self.num_channels, 1) |
|
|
output_conv = nn.Conv2d( |
|
|
self.num_channels, self.num_channels, 3, padding=1) |
|
|
caffe2_xavier_init(lateral_conv) |
|
|
caffe2_xavier_init(output_conv) |
|
|
fpn_laterals.append(lateral_conv) |
|
|
fpn_outputs.append(output_conv) |
|
|
self.fpn_laterals = nn.ModuleList(fpn_laterals) |
|
|
self.fpn_outputs = nn.ModuleList(fpn_outputs) |
|
|
|
|
|
if self.with_ppm: |
|
|
self.ppm = PyramidPoolingModule( |
|
|
self.num_channels, self.num_channels // 4, act_cfg=act_cfg) |
|
|
|
|
|
self.fusion = nn.Conv2d(self.num_channels * 3, self.num_channels, 1) |
|
|
kaiming_init(self.fusion) |
|
|
|
|
|
def forward(self, features): |
|
|
features = features[::-1] |
|
|
prev_features = self.fpn_laterals[0](features[0]) |
|
|
if self.with_ppm: |
|
|
prev_features = self.ppm(prev_features) |
|
|
outputs = [self.fpn_outputs[0](prev_features)] |
|
|
for feature, lat_conv, output_conv in zip(features[1:], |
|
|
self.fpn_laterals[1:], |
|
|
self.fpn_outputs[1:]): |
|
|
lat_features = lat_conv(feature) |
|
|
top_down_features = F.interpolate( |
|
|
prev_features, scale_factor=2.0, mode='nearest') |
|
|
prev_features = lat_features + top_down_features |
|
|
outputs.insert(0, output_conv(prev_features)) |
|
|
size = outputs[0].shape[2:] |
|
|
features = [outputs[0]] + [ |
|
|
F.interpolate(x, size, mode='bilinear', align_corners=False) |
|
|
for x in outputs[1:] |
|
|
] |
|
|
features = self.fusion(torch.cat(features, dim=1)) |
|
|
return features |
|
|
|
