|
| 1 | +"""Adapted from https://github.com/milesial/Pytorch-UNet/tree/master/unet""" |
| 2 | +import torch |
| 3 | +import torch.nn as nn |
| 4 | +import torch.nn.functional as F |
| 5 | + |
| 6 | + |
| 7 | +class UNet(nn.Module): |
| 8 | + def __init__(self, width_multiplier=1, trilinear=True, use_ds_conv=False): |
| 9 | + # Call parent class |
| 10 | + super(UNet, self).__init__() |
| 11 | + |
| 12 | + # Initializations |
| 13 | + _channels = (32, 64, 128, 256, 512) |
| 14 | + conv_type = DepthwiseSeparableConv3d if use_ds_conv else nn.Conv3d |
| 15 | + factor = 2 if trilinear else 1 |
| 16 | + |
| 17 | + # Instance attributes |
| 18 | + self.channels = [int(c * width_multiplier) for c in _channels] |
| 19 | + self.trilinear = trilinear |
| 20 | + |
| 21 | + # Contracting layers |
| 22 | + self.inc = DoubleConv(1, self.channels[0], conv_type=conv_type) |
| 23 | + self.down1 = Down(self.channels[0], self.channels[1], conv_type=conv_type) |
| 24 | + self.down2 = Down(self.channels[1], self.channels[2], conv_type=conv_type) |
| 25 | + self.down3 = Down(self.channels[2], self.channels[3], conv_type=conv_type) |
| 26 | + self.down4 = Down(self.channels[3], self.channels[4] // factor, conv_type=conv_type) |
| 27 | + |
| 28 | + # Expanding layers |
| 29 | + self.up1 = Up(self.channels[4], self.channels[3] // factor, trilinear) |
| 30 | + self.up2 = Up(self.channels[3], self.channels[2] // factor, trilinear) |
| 31 | + self.up3 = Up(self.channels[2], self.channels[1] // factor, trilinear) |
| 32 | + self.up4 = Up(self.channels[1], self.channels[0], trilinear) |
| 33 | + self.outc = OutConv(self.channels[0], 1) |
| 34 | + |
| 35 | + def forward(self, x): |
| 36 | + # Contracting layers |
| 37 | + x1 = self.inc(x) |
| 38 | + x2 = self.down1(x1) |
| 39 | + x3 = self.down2(x2) |
| 40 | + x4 = self.down3(x3) |
| 41 | + x5 = self.down4(x4) |
| 42 | + |
| 43 | + # Expanding layers |
| 44 | + x = self.up1(x5, x4) |
| 45 | + x = self.up2(x, x3) |
| 46 | + x = self.up3(x, x2) |
| 47 | + x = self.up4(x, x1) |
| 48 | + logits = self.outc(x) |
| 49 | + return logits |
| 50 | + |
| 51 | + |
| 52 | +class DoubleConv(nn.Module): |
| 53 | + """(convolution => [BN] => ReLU) * 2""" |
| 54 | + |
| 55 | + def __init__(self, in_channels, out_channels, conv_type=nn.Conv3d, mid_channels=None): |
| 56 | + super().__init__() |
| 57 | + if not mid_channels: |
| 58 | + mid_channels = out_channels |
| 59 | + self.double_conv = nn.Sequential( |
| 60 | + conv_type(in_channels, mid_channels, kernel_size=3, padding=1), |
| 61 | + nn.BatchNorm3d(mid_channels), |
| 62 | + nn.ReLU(inplace=True), |
| 63 | + conv_type(mid_channels, out_channels, kernel_size=3, padding=1), |
| 64 | + nn.BatchNorm3d(out_channels), |
| 65 | + nn.ReLU(inplace=True) |
| 66 | + ) |
| 67 | + |
| 68 | + def forward(self, x): |
| 69 | + return self.double_conv(x) |
| 70 | + |
| 71 | + |
| 72 | +class Down(nn.Module): |
| 73 | + """Downscaling with maxpool then double conv""" |
| 74 | + |
| 75 | + def __init__(self, in_channels, out_channels, conv_type=nn.Conv3d): |
| 76 | + super().__init__() |
| 77 | + self.maxpool_conv = nn.Sequential( |
| 78 | + nn.MaxPool3d(2), |
| 79 | + DoubleConv(in_channels, out_channels, conv_type=conv_type) |
| 80 | + ) |
| 81 | + |
| 82 | + def forward(self, x): |
| 83 | + return self.maxpool_conv(x) |
| 84 | + |
| 85 | + |
| 86 | +class Up(nn.Module): |
| 87 | + """Upscaling then double conv""" |
| 88 | + |
| 89 | + def __init__(self, in_channels, out_channels, trilinear=True): |
| 90 | + super().__init__() |
| 91 | + |
| 92 | + # if trilinear, use the normal convolutions to reduce the number of channels |
| 93 | + if trilinear: |
| 94 | + self.up = nn.Upsample(scale_factor=2, mode='trilinear', align_corners=True) |
| 95 | + self.conv = DoubleConv(in_channels, out_channels, mid_channels=in_channels // 2) |
| 96 | + else: |
| 97 | + self.up = nn.ConvTranspose3d(in_channels, in_channels // 2, kernel_size=2, stride=2) |
| 98 | + self.conv = DoubleConv(in_channels, out_channels) |
| 99 | + |
| 100 | + |
| 101 | + def forward(self, x1, x2): |
| 102 | + x1 = self.up(x1) |
| 103 | + # input is CHW |
| 104 | + diffY = x2.size()[2] - x1.size()[2] |
| 105 | + diffX = x2.size()[3] - x1.size()[3] |
| 106 | + |
| 107 | + x1 = F.pad(x1, [diffX // 2, diffX - diffX // 2, |
| 108 | + diffY // 2, diffY - diffY // 2]) |
| 109 | + # if you have padding issues, see |
| 110 | + # https://github.com/HaiyongJiang/U-Net-Pytorch-Unstructured-Buggy/commit/0e854509c2cea854e247a9c615f175f76fbb2e3a |
| 111 | + # https://github.com/xiaopeng-liao/Pytorch-UNet/commit/8ebac70e633bac59fc22bb5195e513d5832fb3bd |
| 112 | + x = torch.cat([x2, x1], dim=1) |
| 113 | + return self.conv(x) |
| 114 | + |
| 115 | + |
| 116 | +class OutConv(nn.Module): |
| 117 | + def __init__(self, in_channels, out_channels): |
| 118 | + super(OutConv, self).__init__() |
| 119 | + self.conv = nn.Conv3d(in_channels, out_channels, kernel_size=1) |
| 120 | + |
| 121 | + def forward(self, x): |
| 122 | + return self.conv(x) |
| 123 | + |
| 124 | + |
| 125 | +class DepthwiseSeparableConv3d(nn.Module): |
| 126 | + def __init__(self, nin, nout, kernel_size, padding, kernels_per_layer=1): |
| 127 | + super(DepthwiseSeparableConv3d, self).__init__() |
| 128 | + self.depthwise = nn.Conv3d(nin, nin * kernels_per_layer, kernel_size=kernel_size, padding=padding, groups=nin) |
| 129 | + self.pointwise = nn.Conv3d(nin * kernels_per_layer, nout, kernel_size=1) |
| 130 | + |
| 131 | + def forward(self, x): |
| 132 | + out = self.depthwise(x) |
| 133 | + out = self.pointwise(out) |
| 134 | + return out |
0 commit comments