Create unet_model.py

unet_model.py

@@ -0,0 +1,73 @@
+import torch
+import torch.nn as nn
+
+class UNet(nn.Module):
+    def __init__(self):
+        super(UNet, self).__init__()
+
+        def conv_block(in_channels, out_channels):
+            return nn.Sequential(
+                nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1),
+                nn.ReLU(inplace=True),
+                nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1),
+                nn.ReLU(inplace=True),
+            )
+
+        # Encoder
+        self.enc1 = conv_block(3, 64)
+        self.enc2 = conv_block(64, 128)
+        self.enc3 = conv_block(128, 256)
+        self.enc4 = conv_block(256, 512)
+
+        self.pool = nn.MaxPool2d(2)
+
+        # Bottleneck
+        self.bottleneck = conv_block(512, 1024)
+
+        # Decoder
+        self.upconv4 = nn.ConvTranspose2d(1024, 512, kernel_size=2, stride=2)
+        self.dec4 = conv_block(1024, 512)
+
+        self.upconv3 = nn.ConvTranspose2d(512, 256, kernel_size=2, stride=2)
+        self.dec3 = conv_block(512, 256)
+
+        self.upconv2 = nn.ConvTranspose2d(256, 128, kernel_size=2, stride=2)
+        self.dec2 = conv_block(256, 128)
+
+        self.upconv1 = nn.ConvTranspose2d(128, 64, kernel_size=2, stride=2)
+        self.dec1 = conv_block(128, 64)
+
+        self.conv_last = nn.Conv2d(64, 1, kernel_size=1)
+
+    def forward(self, x):
+        c1 = self.enc1(x)
+        p1 = self.pool(c1)
+
+        c2 = self.enc2(p1)
+        p2 = self.pool(c2)
+
+        c3 = self.enc3(p2)
+        p3 = self.pool(c3)
+
+        c4 = self.enc4(p3)
+        p4 = self.pool(c4)
+
+        bottleneck = self.bottleneck(p4)
+
+        u4 = self.upconv4(bottleneck)
+        u4 = torch.cat([u4, c4], dim=1)
+        d4 = self.dec4(u4)
+
+        u3 = self.upconv3(d4)
+        u3 = torch.cat([u3, c3], dim=1)
+        d3 = self.dec3(u3)
+
+        u2 = self.upconv2(d3)
+        u2 = torch.cat([u2, c2], dim=1)
+        d2 = self.dec2(u2)
+
+        u1 = self.upconv1(d2)
+        u1 = torch.cat([u1, c1], dim=1)
+        d1 = self.dec1(u1)
+
+        return torch.sigmoid(self.conv_last(d1))  # sigmoid kept (matches BCELoss training)