2019-10-24 19:37:21 +00:00
|
|
|
""" Parts of the U-Net model """
|
2017-08-16 12:24:29 +00:00
|
|
|
|
|
|
|
import torch
|
|
|
|
import torch.nn as nn
|
|
|
|
import torch.nn.functional as F
|
|
|
|
|
2017-08-17 19:16:19 +00:00
|
|
|
|
2019-10-24 19:37:21 +00:00
|
|
|
class DoubleConv(nn.Module):
|
|
|
|
"""(convolution => [BN] => ReLU) * 2"""
|
|
|
|
|
2020-03-31 03:16:55 +00:00
|
|
|
def __init__(self, in_channels, out_channels, mid_channels=None):
|
2019-10-24 19:37:21 +00:00
|
|
|
super().__init__()
|
2020-03-31 03:16:55 +00:00
|
|
|
if not mid_channels:
|
|
|
|
mid_channels = out_channels
|
2019-10-24 19:37:21 +00:00
|
|
|
self.double_conv = nn.Sequential(
|
2020-03-31 03:16:55 +00:00
|
|
|
nn.Conv2d(in_channels, mid_channels, kernel_size=3, padding=1),
|
|
|
|
nn.BatchNorm2d(mid_channels),
|
2017-08-19 08:59:51 +00:00
|
|
|
nn.ReLU(inplace=True),
|
2020-03-31 03:16:55 +00:00
|
|
|
nn.Conv2d(mid_channels, out_channels, kernel_size=3, padding=1),
|
2019-10-24 19:37:21 +00:00
|
|
|
nn.BatchNorm2d(out_channels),
|
2017-08-19 08:59:51 +00:00
|
|
|
nn.ReLU(inplace=True)
|
2017-08-16 12:24:29 +00:00
|
|
|
)
|
2017-08-17 19:16:19 +00:00
|
|
|
|
2017-08-16 12:24:29 +00:00
|
|
|
def forward(self, x):
|
2019-10-24 19:37:21 +00:00
|
|
|
return self.double_conv(x)
|
2017-08-16 12:24:29 +00:00
|
|
|
|
|
|
|
|
2019-10-24 19:37:21 +00:00
|
|
|
class Down(nn.Module):
|
|
|
|
"""Downscaling with maxpool then double conv"""
|
2017-08-17 19:16:19 +00:00
|
|
|
|
2019-10-24 19:37:21 +00:00
|
|
|
def __init__(self, in_channels, out_channels):
|
|
|
|
super().__init__()
|
|
|
|
self.maxpool_conv = nn.Sequential(
|
2017-08-16 12:24:29 +00:00
|
|
|
nn.MaxPool2d(2),
|
2019-10-24 19:37:21 +00:00
|
|
|
DoubleConv(in_channels, out_channels)
|
2017-08-16 12:24:29 +00:00
|
|
|
)
|
|
|
|
|
|
|
|
def forward(self, x):
|
2019-10-24 19:37:21 +00:00
|
|
|
return self.maxpool_conv(x)
|
|
|
|
|
2017-08-16 12:24:29 +00:00
|
|
|
|
2019-10-24 19:37:21 +00:00
|
|
|
class Up(nn.Module):
|
|
|
|
"""Upscaling then double conv"""
|
2017-08-17 19:16:19 +00:00
|
|
|
|
2019-10-24 19:37:21 +00:00
|
|
|
def __init__(self, in_channels, out_channels, bilinear=True):
|
|
|
|
super().__init__()
|
2017-11-30 05:45:19 +00:00
|
|
|
|
2019-10-24 19:37:21 +00:00
|
|
|
# if bilinear, use the normal convolutions to reduce the number of channels
|
2017-11-30 05:45:19 +00:00
|
|
|
if bilinear:
|
2018-06-08 17:27:32 +00:00
|
|
|
self.up = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)
|
2020-04-17 15:00:11 +00:00
|
|
|
self.conv = DoubleConv(in_channels, out_channels, in_channels // 2)
|
2017-11-30 05:45:19 +00:00
|
|
|
else:
|
2021-08-16 00:53:00 +00:00
|
|
|
self.up = nn.ConvTranspose2d(in_channels, in_channels // 2, kernel_size=2, stride=2)
|
2020-03-31 03:16:55 +00:00
|
|
|
self.conv = DoubleConv(in_channels, out_channels)
|
2017-11-30 05:45:19 +00:00
|
|
|
|
2017-08-16 12:24:29 +00:00
|
|
|
def forward(self, x1, x2):
|
|
|
|
x1 = self.up(x1)
|
2018-11-19 10:51:46 +00:00
|
|
|
# input is CHW
|
2020-04-24 13:57:53 +00:00
|
|
|
diffY = x2.size()[2] - x1.size()[2]
|
|
|
|
diffX = x2.size()[3] - x1.size()[3]
|
2018-11-19 10:51:46 +00:00
|
|
|
|
2019-10-24 19:37:21 +00:00
|
|
|
x1 = F.pad(x1, [diffX // 2, diffX - diffX // 2,
|
|
|
|
diffY // 2, diffY - diffY // 2])
|
|
|
|
# if you have padding issues, see
|
2018-11-19 10:51:46 +00:00
|
|
|
# https://github.com/HaiyongJiang/U-Net-Pytorch-Unstructured-Buggy/commit/0e854509c2cea854e247a9c615f175f76fbb2e3a
|
|
|
|
# https://github.com/xiaopeng-liao/Pytorch-UNet/commit/8ebac70e633bac59fc22bb5195e513d5832fb3bd
|
2017-08-16 12:24:29 +00:00
|
|
|
x = torch.cat([x2, x1], dim=1)
|
2019-10-24 19:37:21 +00:00
|
|
|
return self.conv(x)
|
2017-08-16 12:24:29 +00:00
|
|
|
|
2017-08-17 19:16:19 +00:00
|
|
|
|
2019-10-24 19:37:21 +00:00
|
|
|
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)
|
2017-08-16 12:24:29 +00:00
|
|
|
|
|
|
|
def forward(self, x):
|
2019-10-24 19:37:21 +00:00
|
|
|
return self.conv(x)
|