VisionLanguageGroup/MicroscopyMatching

1""" Parts of the U-Net model """2 3import torch4import torch.nn as nn5import torch.nn.functional as F6 7 8class DoubleConv(nn.Module):9    """(convolution => [BN] => ReLU) * 2"""10 11    def __init__(self, in_channels, out_channels, mid_channels=None):12        super().__init__()13        if not mid_channels:14            mid_channels = out_channels15        self.double_conv = nn.Sequential(16            nn.Conv2d(in_channels, mid_channels, kernel_size=3, padding=1, bias=False),17            nn.BatchNorm2d(mid_channels),18            nn.ReLU(inplace=True),19            nn.Conv2d(mid_channels, out_channels, kernel_size=3, padding=1, bias=False),20            nn.BatchNorm2d(out_channels),21            nn.ReLU(inplace=True)22        )23 24    def forward(self, x):25        return self.double_conv(x)26 27 28class Down(nn.Module):29    """Downscaling with maxpool then double conv"""30 31    def __init__(self, in_channels, out_channels):32        super().__init__()33        self.maxpool_conv = nn.Sequential(34            nn.MaxPool2d(2),35            DoubleConv(in_channels, out_channels)36        )37 38    def forward(self, x):39        return self.maxpool_conv(x)40 41 42class Up(nn.Module):43    """Upscaling then double conv"""44 45    def __init__(self, in_channels, out_channels, bilinear=True):46        super().__init__()47 48        # if bilinear, use the normal convolutions to reduce the number of channels49        if bilinear:50            self.up = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)51            self.conv = DoubleConv(in_channels, out_channels, in_channels // 2)52        else:53            self.up = nn.ConvTranspose2d(in_channels, in_channels // 2, kernel_size=2, stride=2)54            self.conv = DoubleConv(in_channels, out_channels)55 56    def forward(self, x1, x2):57        x1 = self.up(x1)58        # input is CHW59        diffY = x2.size()[2] - x1.size()[2]60        diffX = x2.size()[3] - x1.size()[3]61 62        x1 = F.pad(x1, [diffX // 2, diffX - diffX // 2,63                        diffY // 2, diffY - diffY // 2])64        # if you have padding issues, see65        # https://github.com/HaiyongJiang/U-Net-Pytorch-Unstructured-Buggy/commit/0e854509c2cea854e247a9c615f175f76fbb2e3a66        # https://github.com/xiaopeng-liao/Pytorch-UNet/commit/8ebac70e633bac59fc22bb5195e513d5832fb3bd67        x = torch.cat([x2, x1], dim=1)68        return self.conv(x)69 70 71class OutConv(nn.Module):72    def __init__(self, in_channels, out_channels):73        super(OutConv, self).__init__()74        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size=1)75 76    def forward(self, x):77        return self.conv(x)78