Add model architecture

model.py

@@ -0,0 +1,100 @@
+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)