| | 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 |
