|
|
|
|
|
|
|
|
|
|
|
import torch |
|
|
import torch.nn as nn |
|
|
|
|
|
from scripts.maskscratches import Downsample |
|
|
from scripts.maskscratches import DataParallelWithCallback |
|
|
|
|
|
|
|
|
class UNet(nn.Module): |
|
|
def __init__( |
|
|
self, |
|
|
in_channels=3, |
|
|
out_channels=3, |
|
|
depth=5, |
|
|
conv_num=2, |
|
|
wf=6, |
|
|
padding=True, |
|
|
batch_norm=True, |
|
|
up_mode="upsample", |
|
|
with_tanh=False, |
|
|
sync_bn=True, |
|
|
antialiasing=True, |
|
|
): |
|
|
""" |
|
|
Implementation of |
|
|
U-Net: Convolutional Networks for Biomedical Image Segmentation |
|
|
(Ronneberger et al., 2015) |
|
|
https://arxiv.org/abs/1505.04597 |
|
|
Using the default arguments will yield the exact version used |
|
|
in the original paper |
|
|
Args: |
|
|
in_channels (int): number of input channels |
|
|
out_channels (int): number of output channels |
|
|
depth (int): depth of the network |
|
|
wf (int): number of filters in the first layer is 2**wf |
|
|
padding (bool): if True, apply padding such that the input shape |
|
|
is the same as the output. |
|
|
This may introduce artifacts |
|
|
batch_norm (bool): Use BatchNorm after layers with an |
|
|
activation function |
|
|
up_mode (str): one of 'upconv' or 'upsample'. |
|
|
'upconv' will use transposed convolutions for |
|
|
learned upsampling. |
|
|
'upsample' will use bilinear upsampling. |
|
|
""" |
|
|
|
|
|
super().__init__() |
|
|
assert up_mode in ("upconv", "upsample") |
|
|
self.padding = padding |
|
|
self.depth = depth - 1 |
|
|
prev_channels = in_channels |
|
|
|
|
|
self.first = nn.Sequential( |
|
|
*[nn.ReflectionPad2d(3), nn.Conv2d(in_channels, 2 ** wf, kernel_size=7), nn.LeakyReLU(0.2, True)] |
|
|
) |
|
|
prev_channels = 2 ** wf |
|
|
|
|
|
self.down_path = nn.ModuleList() |
|
|
self.down_sample = nn.ModuleList() |
|
|
for i in range(depth): |
|
|
if antialiasing and depth > 0: |
|
|
self.down_sample.append( |
|
|
nn.Sequential( |
|
|
*[ |
|
|
nn.ReflectionPad2d(1), |
|
|
nn.Conv2d(prev_channels, prev_channels, kernel_size=3, stride=1, padding=0), |
|
|
nn.BatchNorm2d(prev_channels), |
|
|
nn.LeakyReLU(0.2, True), |
|
|
Downsample(channels=prev_channels, stride=2), |
|
|
] |
|
|
) |
|
|
) |
|
|
else: |
|
|
self.down_sample.append( |
|
|
nn.Sequential( |
|
|
*[ |
|
|
nn.ReflectionPad2d(1), |
|
|
nn.Conv2d(prev_channels, prev_channels, kernel_size=4, stride=2, padding=0), |
|
|
nn.BatchNorm2d(prev_channels), |
|
|
nn.LeakyReLU(0.2, True), |
|
|
] |
|
|
) |
|
|
) |
|
|
self.down_path.append( |
|
|
UNetConvBlock(conv_num, prev_channels, 2 ** (wf + i + 1), padding, batch_norm) |
|
|
) |
|
|
prev_channels = 2 ** (wf + i + 1) |
|
|
|
|
|
self.up_path = nn.ModuleList() |
|
|
for i in reversed(range(depth)): |
|
|
self.up_path.append( |
|
|
UNetUpBlock(conv_num, prev_channels, 2 ** (wf + i), up_mode, padding, batch_norm) |
|
|
) |
|
|
prev_channels = 2 ** (wf + i) |
|
|
|
|
|
if with_tanh: |
|
|
self.last = nn.Sequential( |
|
|
*[nn.ReflectionPad2d(1), nn.Conv2d(prev_channels, out_channels, kernel_size=3), nn.Tanh()] |
|
|
) |
|
|
else: |
|
|
self.last = nn.Sequential( |
|
|
*[nn.ReflectionPad2d(1), nn.Conv2d(prev_channels, out_channels, kernel_size=3)] |
|
|
) |
|
|
|
|
|
if sync_bn: |
|
|
self = DataParallelWithCallback(self) |
|
|
|
|
|
def forward(self, x): |
|
|
x = self.first(x) |
|
|
|
|
|
blocks = [] |
|
|
for i, down_block in enumerate(self.down_path): |
|
|
blocks.append(x) |
|
|
x = self.down_sample[i](x) |
|
|
x = down_block(x) |
|
|
|
|
|
for i, up in enumerate(self.up_path): |
|
|
x = up(x, blocks[-i - 1]) |
|
|
|
|
|
return self.last(x) |
|
|
|
|
|
|
|
|
class UNetConvBlock(nn.Module): |
|
|
def __init__(self, conv_num, in_size, out_size, padding, batch_norm): |
|
|
super(UNetConvBlock, self).__init__() |
|
|
block = [] |
|
|
|
|
|
for _ in range(conv_num): |
|
|
block.append(nn.ReflectionPad2d(padding=int(padding))) |
|
|
block.append(nn.Conv2d(in_size, out_size, kernel_size=3, padding=0)) |
|
|
if batch_norm: |
|
|
block.append(nn.BatchNorm2d(out_size)) |
|
|
block.append(nn.LeakyReLU(0.2, True)) |
|
|
in_size = out_size |
|
|
|
|
|
self.block = nn.Sequential(*block) |
|
|
|
|
|
def forward(self, x): |
|
|
out = self.block(x) |
|
|
return out |
|
|
|
|
|
|
|
|
class UNetUpBlock(nn.Module): |
|
|
def __init__(self, conv_num, in_size, out_size, up_mode, padding, batch_norm): |
|
|
super(UNetUpBlock, self).__init__() |
|
|
if up_mode == "upconv": |
|
|
self.up = nn.ConvTranspose2d(in_size, out_size, kernel_size=2, stride=2) |
|
|
elif up_mode == "upsample": |
|
|
self.up = nn.Sequential( |
|
|
nn.Upsample(mode="bilinear", scale_factor=2, align_corners=False), |
|
|
nn.ReflectionPad2d(1), |
|
|
nn.Conv2d(in_size, out_size, kernel_size=3, padding=0), |
|
|
) |
|
|
|
|
|
self.conv_block = UNetConvBlock(conv_num, in_size, out_size, padding, batch_norm) |
|
|
|
|
|
def center_crop(self, layer, target_size): |
|
|
_, _, layer_height, layer_width = layer.size() |
|
|
diff_y = (layer_height - target_size[0]) // 2 |
|
|
diff_x = (layer_width - target_size[1]) // 2 |
|
|
return layer[:, :, diff_y: (diff_y + target_size[0]), diff_x: (diff_x + target_size[1])] |
|
|
|
|
|
def forward(self, x, bridge): |
|
|
up = self.up(x) |
|
|
crop1 = self.center_crop(bridge, up.shape[2:]) |
|
|
out = torch.cat([up, crop1], 1) |
|
|
out = self.conv_block(out) |
|
|
|
|
|
return out |
|
|
|
|
|
|
|
|
|
