|
|
import torch |
|
|
from torch import nn |
|
|
from torch.nn import functional as F |
|
|
|
|
|
from .utils import _SimpleSegmentationModel |
|
|
from ._deeplab import ASPPConv, ASPPPooling, ASPP, AtrousSeparableConvolution |
|
|
from .enhanced_modules import EOANetModule |
|
|
|
|
|
|
|
|
class EnhancedDeepLabV3(_SimpleSegmentationModel): |
|
|
""" |
|
|
Implements Enhanced DeepLabV3 model with Normalized Multi-Scale Attention and Entropy-Optimized Gating. |
|
|
""" |
|
|
pass |
|
|
|
|
|
|
|
|
class EnhancedDeepLabHeadV3Plus(nn.Module): |
|
|
def __init__(self, in_channels, low_level_channels, num_classes, aspp_dilate=[12, 24, 36], |
|
|
use_eoaNet=True, msa_scales=[1, 2, 4], eog_beta=0.5): |
|
|
super(EnhancedDeepLabHeadV3Plus, self).__init__() |
|
|
self.use_eoaNet = use_eoaNet |
|
|
|
|
|
self.project = nn.Sequential( |
|
|
nn.Conv2d(low_level_channels, 48, 1, bias=False), |
|
|
nn.BatchNorm2d(48), |
|
|
nn.ReLU(inplace=True), |
|
|
) |
|
|
|
|
|
self.aspp = ASPP(in_channels, aspp_dilate) |
|
|
|
|
|
|
|
|
if self.use_eoaNet: |
|
|
self.eoaNet = EOANetModule(256, scales=msa_scales, beta=eog_beta) |
|
|
|
|
|
self.classifier = nn.Sequential( |
|
|
nn.Conv2d(304, 256, 3, padding=1, bias=False), |
|
|
nn.BatchNorm2d(256), |
|
|
nn.ReLU(inplace=True), |
|
|
nn.Conv2d(256, num_classes, 1) |
|
|
) |
|
|
self._init_weight() |
|
|
|
|
|
def forward(self, feature): |
|
|
low_level_feature = self.project(feature['low_level']) |
|
|
output_feature = self.aspp(feature['out']) |
|
|
|
|
|
|
|
|
if self.use_eoaNet: |
|
|
output_feature = self.eoaNet(output_feature) |
|
|
|
|
|
output_feature = F.interpolate(output_feature, size=low_level_feature.shape[2:], mode='bilinear', align_corners=False) |
|
|
return self.classifier(torch.cat([low_level_feature, output_feature], dim=1)) |
|
|
|
|
|
def _init_weight(self): |
|
|
for m in self.modules(): |
|
|
if isinstance(m, nn.Conv2d): |
|
|
nn.init.kaiming_normal_(m.weight) |
|
|
elif isinstance(m, (nn.BatchNorm2d, nn.GroupNorm)): |
|
|
nn.init.constant_(m.weight, 1) |
|
|
nn.init.constant_(m.bias, 0) |
|
|
|
|
|
|
|
|
class EnhancedDeepLabHead(nn.Module): |
|
|
def __init__(self, in_channels, num_classes, aspp_dilate=[12, 24, 36], |
|
|
use_eoaNet=True, msa_scales=[1, 2, 4], eog_beta=0.5): |
|
|
super(EnhancedDeepLabHead, self).__init__() |
|
|
self.use_eoaNet = use_eoaNet |
|
|
|
|
|
self.aspp = ASPP(in_channels, aspp_dilate) |
|
|
|
|
|
|
|
|
if self.use_eoaNet: |
|
|
self.eoaNet = EOANetModule(256, scales=msa_scales, beta=eog_beta) |
|
|
|
|
|
self.classifier = nn.Sequential( |
|
|
nn.Conv2d(256, 256, 3, padding=1, bias=False), |
|
|
nn.BatchNorm2d(256), |
|
|
nn.ReLU(inplace=True), |
|
|
nn.Conv2d(256, num_classes, 1) |
|
|
) |
|
|
self._init_weight() |
|
|
|
|
|
def forward(self, feature): |
|
|
output = self.aspp(feature['out']) |
|
|
|
|
|
|
|
|
if self.use_eoaNet: |
|
|
output = self.eoaNet(output) |
|
|
|
|
|
return self.classifier(output) |
|
|
|
|
|
def _init_weight(self): |
|
|
for m in self.modules(): |
|
|
if isinstance(m, nn.Conv2d): |
|
|
nn.init.kaiming_normal_(m.weight) |
|
|
elif isinstance(m, (nn.BatchNorm2d, nn.GroupNorm)): |
|
|
nn.init.constant_(m.weight, 1) |
|
|
nn.init.constant_(m.bias, 0) |
|
|
|
|
|
|
|
|
def convert_to_separable_conv(module): |
|
|
new_module = module |
|
|
if isinstance(module, nn.Conv2d) and module.kernel_size[0]>1: |
|
|
new_module = AtrousSeparableConvolution(module.in_channels, |
|
|
module.out_channels, |
|
|
module.kernel_size, |
|
|
module.stride, |
|
|
module.padding, |
|
|
module.dilation, |
|
|
module.bias) |
|
|
for name, child in module.named_children(): |
|
|
new_module.add_module(name, convert_to_separable_conv(child)) |
|
|
return new_module |
