| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
|
|
| class DoubleConv(nn.Module): |
| def __init__(self, in_channels, out_channels, dropout=0.1): |
| super().__init__() |
| self.conv = nn.Sequential( |
| nn.Conv2d(in_channels, out_channels, 3, 1, 1, bias=True), |
| nn.BatchNorm2d(out_channels), |
| nn.ReLU(inplace=True), |
| nn.Dropout2d(dropout), |
| nn.Conv2d(out_channels, out_channels, 3, 1, 1, bias=True), |
| nn.BatchNorm2d(out_channels), |
| nn.ReLU(inplace=True), |
| ) |
|
|
| def forward(self, x): |
| return self.conv(x) |
|
|
| class AttentionGate(nn.Module): |
| def __init__(self, F_g, F_l, F_int): |
| super().__init__() |
| self.W_g = nn.Conv2d(F_g, F_int, kernel_size=1, stride=1, padding=0, bias=True) |
| self.W_x = nn.Conv2d(F_l, F_int, kernel_size=1, stride=1, padding=0, bias=True) |
| self.psi = nn.Sequential( |
| nn.Conv2d(F_int, 1, kernel_size=1, stride=1, padding=0, bias=True), |
| nn.Sigmoid() |
| ) |
| self.relu = nn.ReLU(inplace=True) |
|
|
| def forward(self, g, x): |
| g1 = self.W_g(g) |
| x1 = self.W_x(x) |
| psi = self.relu(g1 + x1) |
| psi = self.psi(psi) |
| return x * psi |
|
|
| class AttentionUNet(nn.Module): |
| def __init__(self, img_ch=1, output_ch=4): |
| super().__init__() |
| self.Maxpool = nn.MaxPool2d(kernel_size=2, stride=2) |
|
|
| self.downs = nn.ModuleList([ |
| DoubleConv(img_ch, 64), |
| DoubleConv(64, 128), |
| DoubleConv(128, 256), |
| DoubleConv(256, 512) |
| ]) |
|
|
| self.bottleneck = DoubleConv(512, 1024) |
|
|
| self.ups = nn.ModuleList([ |
| nn.ConvTranspose2d(1024, 512, kernel_size=2, stride=2), |
| nn.ConvTranspose2d(512, 256, kernel_size=2, stride=2), |
| nn.ConvTranspose2d(256, 128, kernel_size=2, stride=2), |
| nn.ConvTranspose2d(128, 64, kernel_size=2, stride=2) |
| ]) |
| |
| self.attention_gates = nn.ModuleList([ |
| AttentionGate(F_g=512, F_l=512, F_int=256), |
| AttentionGate(F_g=256, F_l=256, F_int=128), |
| AttentionGate(F_g=128, F_l=128, F_int=64), |
| AttentionGate(F_g=64, F_l=64, F_int=32) |
| ]) |
| |
| self.up_convs = nn.ModuleList([ |
| DoubleConv(1024, 512), |
| DoubleConv(512, 256), |
| DoubleConv(256, 128), |
| DoubleConv(128, 64) |
| ]) |
|
|
| self.final_conv = nn.Conv2d(64, output_ch, kernel_size=1, stride=1, padding=0) |
|
|
| def forward(self, x): |
| e1 = self.downs[0](x) |
| e2 = self.downs[1](self.Maxpool(e1)) |
| e3 = self.downs[2](self.Maxpool(e2)) |
| e4 = self.downs[3](self.Maxpool(e3)) |
|
|
| b = self.bottleneck(self.Maxpool(e4)) |
|
|
| d4 = self.ups[0](b) |
| x4 = self.attention_gates[0](g=d4, x=e4) |
| d4 = self.up_convs[0](torch.cat((x4, d4), dim=1)) |
|
|
| d3 = self.ups[1](d4) |
| x3 = self.attention_gates[1](g=d3, x=e3) |
| d3 = self.up_convs[1](torch.cat((x3, d3), dim=1)) |
|
|
| d2 = self.ups[2](d3) |
| x2 = self.attention_gates[2](g=d2, x=e2) |
| d2 = self.up_convs[2](torch.cat((x2, d2), dim=1)) |
|
|
| d1 = self.ups[3](d2) |
| x1 = self.attention_gates[3](g=d1, x=e1) |
| d1 = self.up_convs[3](torch.cat((x1, d1), dim=1)) |
|
|
| return self.final_conv(d1) |
|
|
