|
import math |
|
|
|
import torch |
|
import torch.nn as nn |
|
import torch.nn.functional as F |
|
|
|
|
|
class SqeezeExcite(nn.Module): |
|
def __init__(self, channel, reduction_ratio = 16): |
|
super(SqeezeExcite,self).__init__() |
|
self.GAP = nn.AdaptiveAvgPool2d(1) |
|
|
|
self.mlp = nn.Sequential( |
|
nn.Linear(channel, channel//reduction_ratio, bias = False), |
|
nn.ReLU(inplace=True), |
|
nn.Linear(channel//reduction_ratio,channel,bias = False), |
|
nn.Sigmoid() |
|
) |
|
|
|
def forward(self,x): |
|
b,c,_,_ = x.size() |
|
out = self.GAP(x).view(b,c) |
|
out = self.mlp(out).view(b,c,1,1) |
|
return x * out.expand_as(x) |
|
|
|
|
|
class ECA(nn.Module): |
|
|
|
def __init__(self,channels, b = 1, gamma = 2): |
|
super(ECA, self).__init__() |
|
self.GAP = nn.AdaptiveAvgPool2d(1) |
|
self.channels = channels |
|
self.b = b |
|
self.gamma = gamma |
|
self.conv = nn.Conv1d(1, 1, kernel_size=self.adaptive_kernel(),padding = (self.adaptive_kernel()-1)//2, bias = False) |
|
self.sigmoid = nn.Sigmoid() |
|
|
|
def forward(self,x): |
|
attn = self.GAP(x) |
|
attn = self.conv(attn.squeeze(-1).transpose(-1,-2)).transpose(-1,-2).unsqueeze(-1) |
|
attn = self.sigmoid(attn) |
|
return x * attn.expand_as(x) |
|
|
|
|
|
def adaptive_kernel(self): |
|
k = int(abs(math.log2(self.channels)/self.gamma) + self.b) |
|
ksize = k if k%2 else k+1 |
|
return ksize |
|
|
|
|
|
|
|
class UNetConvBlock(nn.Module): |
|
def __init__(self, in_channel, out_channel, ca_layer): |
|
super(UNetConvBlock, self).__init__() |
|
block = [] |
|
|
|
block.append(nn.Conv2d(in_channel, out_channel, kernel_size=3, stride=1, padding=1)) |
|
block.append(nn.PReLU()) |
|
|
|
block.append(nn.Conv2d(out_channel, out_channel, kernel_size=3, padding=1, stride=1)) |
|
block.append(nn.PReLU()) |
|
|
|
if ca_layer: |
|
block.append(ECA(out_channel)) |
|
self.block = nn.Sequential(*block) |
|
|
|
def forward(self, x): |
|
out = self.block(x) |
|
return out |
|
|
|
class AttentionGate(nn.Module): |
|
def __init__(self, F_g, F_l, dimensions): |
|
super(AttentionGate, self).__init__() |
|
self.W_gate = nn.Sequential( |
|
nn.Conv2d(F_g, dimensions, kernel_size=1, stride=1, padding=0, bias=True), |
|
nn.BatchNorm2d(dimensions) |
|
) |
|
|
|
self.W_x = nn.Sequential( |
|
nn.Conv2d(F_l, dimensions, kernel_size=1, stride=1, padding=0, bias=True), |
|
nn.BatchNorm2d(dimensions) |
|
) |
|
|
|
self.psi = nn.Sequential( |
|
nn.Conv2d(dimensions, 1, kernel_size=1, stride=1, padding=0, bias=True), |
|
nn.BatchNorm2d(1), |
|
nn.Sigmoid() |
|
) |
|
|
|
self.relu = nn.PReLU() |
|
|
|
def forward(self, g, x): |
|
g1 = self.W_gate(g) |
|
x1 = self.W_x(x) |
|
psi = self.relu(g1 + x1) |
|
psi = self.psi(psi) |
|
out = x * psi |
|
return out |
|
|
|
class UNetUpConvBlock(nn.Module): |
|
def __init__(self, in_channel, out_channel, upmode, ca_layer, up_factor = 2, att_mode = "standard"): |
|
super(UNetUpConvBlock, self).__init__() |
|
self.att_mode = att_mode |
|
self.ca_layer = ca_layer |
|
if upmode == 'upsample': |
|
self.Upsize = nn.Sequential( |
|
nn.Upsample(scale_factor=up_factor, mode='bilinear', align_corners=False), |
|
nn.Conv2d(in_channel, out_channel, kernel_size=1, stride=1, padding=0), |
|
) |
|
elif upmode == 'upconv': |
|
self.Upsize = nn.ConvTranspose2d(in_channel,out_channel,kernel_size=2,stride = 2) |
|
elif upmode == 'shuffle': |
|
self.Upsize = nn.Sequential( |
|
nn.Conv2d(in_channel,out_channel*4,kernel_size=3,stride=1,padding=1), |
|
nn.PReLU(), |
|
nn.PixelShuffle(2), |
|
nn.Conv2d(out_channel,out_channel,kernel_size=3,stride = 1,padding=1) |
|
) |
|
|
|
|
|
|
|
if self.att_mode == 'standard': |
|
self.attention_gate = AttentionGate(out_channel, out_channel, out_channel) |
|
self.conv = UNetConvBlock(in_channel, out_channel, ca_layer=self.ca_layer) |
|
elif self.att_mode == 'modified': |
|
self.attention_gate = AttentionGate(out_channel, out_channel, out_channel ) |
|
self.conv = UNetConvBlock(3*out_channel, out_channel, ca_layer = self.ca_layer) |
|
elif self.att_mode == 'None': |
|
self.conv = UNetConvBlock(in_channel, out_channel, ca_layer=self.ca_layer) |
|
|
|
|
|
|
|
def forward(self, x, residue): |
|
x = self.Upsize(x) |
|
x = F.interpolate(x, size=residue.shape[2:], mode='bilinear') |
|
if self.att_mode == "standard": |
|
attn = self.attention_gate(g = x, x=residue) |
|
out = torch.cat([x, attn],dim = 1) |
|
out = self.conv(out) |
|
elif self.att_mode == 'modified': |
|
attn = self.attention_gate(g = x, x = residue) |
|
out = torch.cat([x,residue,attn],dim = 1) |
|
out = self.conv(out) |
|
elif self.att_mode == 'None': |
|
out = torch.cat([x,residue], dim = 1) |
|
out = self.conv(out) |
|
return out |
|
|
|
|
|
|
|
|
|
|
|
class AUNet(nn.Module): |
|
def __init__(self,in_channels = 6,out_channels = 6,depth = 3,growth_factor = 6, |
|
interp_mode = 'bicubic', up_mode = 'upconv',spatial_attention = "standard", ca_layer = True): |
|
super(AUNet,self).__init__() |
|
|
|
if not spatial_attention in ['None', 'modified', 'standard']: |
|
raise AssertionError("spatial_attention options : \'None\'- no spatial attention, \'standard\'-spatial attention as in attention unet paper, \'modified\'-modified attention unet") |
|
|
|
self.in_channels = in_channels |
|
self.out_channels = out_channels |
|
self.depth = depth |
|
self.growth_factor = growth_factor |
|
self.interp_mode = interp_mode |
|
prev_channels = self.in_channels |
|
self.up_mode = up_mode |
|
self.att_mode = spatial_attention |
|
self.ca_layer = ca_layer |
|
|
|
self.encoding_module = nn.ModuleList() |
|
for i in range(self.depth): |
|
self.encoding_module.append(UNetConvBlock(in_channel=prev_channels,out_channel=2**(self.growth_factor + i), ca_layer=self.ca_layer)) |
|
prev_channels = 2**(self.growth_factor+i) |
|
|
|
self.decoding_module = nn.ModuleList() |
|
for i in reversed(range(self.depth-1)): |
|
self.decoding_module.append(UNetUpConvBlock(prev_channels,2**(self.growth_factor+i),upmode = self.up_mode, att_mode = self.att_mode, ca_layer = self.ca_layer)) |
|
prev_channels = 2**(self.growth_factor+i) |
|
|
|
self.final = nn.Conv2d(prev_channels,out_channels,1,1,0) |
|
|
|
def forward(self,MS,PAN = None): |
|
if PAN == None: |
|
x = MS |
|
else: |
|
x = torch.cat([MS,PAN],dim = 1) |
|
blocks = [] |
|
for i,down in enumerate(self.encoding_module): |
|
x = down(x) |
|
if i != len(self.encoding_module)-1: |
|
blocks.append(x) |
|
x = F.avg_pool2d(x,2) |
|
|
|
for i,up in enumerate(self.decoding_module): |
|
x = up(x,blocks[-i-1]) |
|
|
|
x = self.final(x) |
|
return x |
|
|
|
|
|
if __name__ == '__main__': |
|
x = torch.rand([9,7,256,256]).cuda() |
|
model = AUNet(in_channels=7, out_channels=6, depth=5, spatial_attention="modified", growth_factor=6, |
|
interp_mode='bilinear', up_mode='upconv', ca_layer=True).cuda() |
|
x = model(x) |
|
|
|
|
|
|
|
activation = {} |
|
for layer in model: |
|
print(layer) |
