unet.py

""" Parts of the U-Net model """

import torch
import torch.nn as nn
import torch.nn.functional as F


class UNet(nn.Module):
    def __init__(self, n_channels=3, N=256):
        super(UNet, self).__init__()
        self.down1 = DoubleConv(n_channels, N // 4)
        self.down2 = DoubleConv(N // 4, N // 2)
        self.down3 = DoubleConv(N // 2, N)
        self.down4 = DoubleConv(N, 2 * N)

        self.inter = nn.Sequential(
            nn.Conv2d(2 * N, 4 * N, kernel_size=3, stride=1, padding=1),
            nn.LeakyReLU(inplace=True),
            nn.Conv2d(4 * N, 4 * N, kernel_size=3, stride=1, padding=1),
            nn.LeakyReLU(inplace=True),
        )

        self.up1 = Up(4 * N, 2 * N)
        self.up2 = Up(2 * N, N)
        self.up3 = Up(N, N // 2)
        self.up4 = Up(N // 2, N // 4)
        self.out = nn.Conv2d(N // 4, n_channels, kernel_size=3, stride=1, padding=1)

        self.maxpool = nn.MaxPool2d(2)

    def forward(self, x):
        identity = x
        x1 = self.down1(x)
        x1_pooled = self.maxpool(x1)

        x2 = self.down2(x1_pooled)
        x2_pooled = self.maxpool(x2)

        x3 = self.down3(x2_pooled)
        x3_pooled = self.maxpool(x3)

        x4 = self.down4(x3_pooled)
        x4_pooled = self.maxpool(x4)

        x5 = self.inter(x4_pooled)

        x = self.up1(x5, x4)
        x = self.up2(x, x3)
        x = self.up3(x, x2)
        x = self.up4(x, x1)
        x = self.out(x)

        return x + identity


class DoubleConv(nn.Module):
    """(convolution => [BN] => ReLU) * 2"""

    def __init__(self, in_channels, out_channels, mid_channels=None):
        super().__init__()
        if not mid_channels:
            mid_channels = out_channels
        self.double_conv = nn.Sequential(
            nn.Conv2d(in_channels, mid_channels, kernel_size=3, padding=1, bias=False),
            nn.LeakyReLU(inplace=True),
            nn.Conv2d(mid_channels, out_channels, kernel_size=3, padding=1, bias=False),
            nn.LeakyReLU(inplace=True)
        )

    def forward(self, x):
        return self.double_conv(x)


class Down(nn.Module):
    """Downscaling with maxpool then double conv"""

    def __init__(self, in_channels, out_channels):
        super().__init__()
        self.maxpool_conv = nn.Sequential(
            nn.MaxPool2d(2),
            DoubleConv(in_channels, out_channels)
        )

    def forward(self, x):
        return self.maxpool_conv(x)


class Up(nn.Module):
    """Upscaling then double conv"""

    def __init__(self, in_channels, out_channels):
        super().__init__()
        self.up = nn.ConvTranspose2d(in_channels, in_channels // 2, kernel_size=3, stride=2, padding=1, output_padding=1)
        self.conv = DoubleConv(in_channels, out_channels)

    def forward(self, x1, x2):
        x1 = self.up(x1)
        x = torch.cat([x2, x1], dim=1)
        return self.conv(x)


class OutConv(nn.Module):
    def __init__(self, in_channels, out_channels):
        super(OutConv, self).__init__()
        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size=1)

    def forward(self, x):
        return self.conv(x)