def __init__(
self, input_size: Tuple[int, int, int] = (1, 128, 128)) -> None:
super().__init__()
self.input_size = input_size
self.enc1 = UNetEncoderBlock(in_channels=input_size[0],
out_channels=64)
self.enc2 = UNetEncoderBlock(in_channels=64, out_channels=128)
self.enc3 = UNetEncoderBlock(in_channels=128, out_channels=256)
self.enc4 = UNetEncoderBlock(in_channels=256, out_channels=512)
self.bottle_neck = UNetDecoderBlock(in_channels=512,
mid_channels=1024,
out_channels=512)
self.dec4 = UNetDecoderBlock(in_channels=1024,
mid_channels=512,
out_channels=256)
self.dec3 = UNetDecoderBlock(in_channels=512,
mid_channels=256,
out_channels=128)
self.dec2 = UNetDecoderBlock(in_channels=256,
mid_channels=128,
out_channels=64)
self.dec1 = nn.Sequential(nn.Conv2d(128, 64, 3, padding=1),
nn.ReLU(inplace=True),
nn.Conv2d(64, 64, 3, padding=1),
nn.ReLU(inplace=True), nn.Conv2d(64, 1, 1),
nn.Sigmoid())
for layer in self.dec1:
if isinstance(layer, nn.Conv2d):
he_normal(layer.weight.data)
layer.bias.data.zero_()
def __init__(self, in_channels, out_channels):
super().__init__()
self.upsample = nn.Sequential(
nn.Upsample(mode="bilinear", scale_factor=2, align_corners=True),
nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1),
nn.ReLU(inplace=True))
for l in self.upsample:
if isinstance(l, nn.Conv2d):
he_normal(l.weight.data)
def __init__(self, in_channels: int, out_channels: int) -> None:
super().__init__()
self.layers = nn.Sequential(nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1),
nn.ReLU(inplace=True),
nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1),
nn.ReLU(inplace=True))
for layer in self.layers:
if isinstance(layer, nn.Conv2d):
he_normal(layer.weight.data)
layer.bias.data.zero_()
def __init__(self, upsample_in_ch, conv_in_ch, out_ch):
super().__init__()
self.upsample = Upsample2D(upsample_in_ch, out_ch)
self.conv_layers = nn.Sequential(
nn.Conv2d(conv_in_ch, out_ch, kernel_size=3, padding=1),
nn.ReLU(inplace=True),
nn.Conv2d(out_ch, out_ch, kernel_size=3, padding=1),
nn.ReLU(inplace=True))
for l in self.conv_layers:
if isinstance(l, nn.Conv2d):
he_normal(l.weight.data)
def __init__(self, in_channels: int, mid_channels: int, out_channels: int) -> None:
super().__init__()
self.layers = nn.Sequential(nn.Conv2d(in_channels, mid_channels, 3, padding=1),
nn.ReLU(inplace=True),
nn.Conv2d(mid_channels, mid_channels, 3, padding=1),
nn.ReLU(inplace=True),
nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True),
nn.Conv2d(mid_channels, out_channels, 3, padding=1),
nn.ReLU(inplace=True))
for layer in self.layers:
if isinstance(layer, nn.Conv2d):
he_normal(layer.weight.data)
layer.bias.data.zero_()
