def __init__(self, in_channels, out_channels=3, output_heads=1, conv_layer_order='crg', init_channel_number=32,
**kwargs):
super(TagsUNet3D, self).__init__()
# number of groups for the GroupNorm
num_groups = min(init_channel_number // 2, 32)
# encoder path consist of 4 subsequent Encoder modules
# the number of features maps is the same as in the paper
self.encoders = nn.ModuleList([
Encoder(in_channels, init_channel_number, apply_pooling=False, conv_layer_order=conv_layer_order,
num_groups=num_groups),
Encoder(init_channel_number, 2 * init_channel_number, conv_layer_order=conv_layer_order,
num_groups=num_groups),
Encoder(2 * init_channel_number, 4 * init_channel_number, conv_layer_order=conv_layer_order,
num_groups=num_groups),
Encoder(4 * init_channel_number, 8 * init_channel_number, conv_layer_order=conv_layer_order,
num_groups=num_groups)
])
self.decoders = nn.ModuleList([
Decoder(4 * init_channel_number + 8 * init_channel_number, 4 * init_channel_number,
conv_layer_order=conv_layer_order, num_groups=num_groups),
Decoder(2 * init_channel_number + 4 * init_channel_number, 2 * init_channel_number,
conv_layer_order=conv_layer_order, num_groups=num_groups),
Decoder(init_channel_number + 2 * init_channel_number, init_channel_number,
conv_layer_order=conv_layer_order, num_groups=num_groups)
])
self.final_heads = nn.ModuleList(
[FinalConv(init_channel_number, out_channels, num_groups=num_groups) for _ in
range(output_heads)])
def __init__(self, in_channels, out_channels, final_sigmoid, init_channel_number=32, **kwargs):
super(DistanceTransformUNet3D, self).__init__()
# number of groups for the GroupNorm
num_groups = min(init_channel_number // 2, 32)
# encoder path consist of 4 subsequent Encoder modules
# the number of features maps is the same as in the paper
self.encoders = nn.ModuleList([
Encoder(in_channels, init_channel_number, apply_pooling=False, conv_layer_order='crg',
num_groups=num_groups),
Encoder(init_channel_number, 2 * init_channel_number, pool_type='avg', conv_layer_order='crg',
num_groups=num_groups)
])
self.decoders = nn.ModuleList([
Decoder(3 * init_channel_number, init_channel_number, conv_layer_order='crg', num_groups=num_groups)
])
# in the last layer a 1×1 convolution reduces the number of output
# channels to the number of labels
self.final_conv = nn.Conv3d(init_channel_number, out_channels, 1)
if final_sigmoid:
self.final_activation = nn.Sigmoid()
else:
self.final_activation = nn.Softmax(dim=1)
def __init__(self,
in_channels,
out_channels,
final_sigmoid,
f_maps=64,
layer_order='crg',
num_groups=8,
**kwargs):
super(UNet3D, self).__init__()
if isinstance(f_maps, int):
# use 4 levels in the encoder path as suggested in the paper
# f_maps = [f_maps, f_maps*2, f_maps*4, f_maps*8]
f_maps = create_feature_maps(f_maps, number_of_fmaps=4)
# create encoder path consisting of Encoder modules. The length of the encoder is equal to `len(f_maps)`
# uses DoubleConv as a basic_module for the Encoder
encoders = []
for i, out_feature_num in enumerate(f_maps):
if i == 0:
encoder = Encoder(in_channels,
out_feature_num,
apply_pooling=False,
basic_module=DoubleConv,
conv_layer_order=layer_order,
num_groups=num_groups)
else:
encoder = Encoder(f_maps[i - 1],
out_feature_num,
basic_module=DoubleConv,
conv_layer_order=layer_order,
num_groups=num_groups)
encoders.append(encoder)
self.encoders = nn.ModuleList(encoders)
# create decoder path consisting of the Decoder modules. The length of the decoder is equal to `len(f_maps) - 1`
# uses DoubleConv as a basic_module for the Decoder
decoders = []
reversed_f_maps = list(reversed(f_maps))
for i in range(len(reversed_f_maps) - 1):
in_feature_num = reversed_f_maps[i] + reversed_f_maps[i + 1]
out_feature_num = reversed_f_maps[i + 1]
decoder = Decoder(in_feature_num,
out_feature_num,
basic_module=DoubleConv,
conv_layer_order=layer_order,
num_groups=num_groups)
decoders.append(decoder)
self.decoders = nn.ModuleList(decoders)
# in the last layer a 1×1 convolution reduces the number of output
# channels to the number of labels
self.final_conv = nn.Conv3d(f_maps[0], out_channels, 1)
if final_sigmoid:
self.final_activation = nn.Sigmoid()
else:
self.final_activation = nn.Softmax(dim=1)
def __init__(self,
in_channels,
out_channels,
f_maps=16,
num_groups=8,
**kwargs):
super(Noise2NoiseUNet3D, self).__init__()
# Use LeakyReLU activation everywhere except the last layer
conv_layer_order = 'clg'
if isinstance(f_maps, int):
# use 5 levels in the encoder path as suggested in the paper
f_maps = create_feature_maps(f_maps, number_of_fmaps=5)
# create encoder path consisting of Encoder modules. The length of the encoder is equal to `len(f_maps)`
# uses DoubleConv as a basic_module for the Encoder
encoders = []
for i, out_feature_num in enumerate(f_maps):
if i == 0:
encoder = Encoder(in_channels,
out_feature_num,
apply_pooling=False,
basic_module=DoubleConv,
conv_layer_order=conv_layer_order,
num_groups=num_groups)
else:
encoder = Encoder(f_maps[i - 1],
out_feature_num,
basic_module=DoubleConv,
conv_layer_order=conv_layer_order,
num_groups=num_groups)
encoders.append(encoder)
self.encoders = nn.ModuleList(encoders)
# create decoder path consisting of the Decoder modules. The length of the decoder is equal to `len(f_maps) - 1`
# uses DoubleConv as a basic_module for the Decoder
decoders = []
reversed_f_maps = list(reversed(f_maps))
for i in range(len(reversed_f_maps) - 1):
in_feature_num = reversed_f_maps[i] + reversed_f_maps[i + 1]
out_feature_num = reversed_f_maps[i + 1]
decoder = Decoder(in_feature_num,
out_feature_num,
basic_module=DoubleConv,
conv_layer_order=conv_layer_order,
num_groups=num_groups)
decoders.append(decoder)
self.decoders = nn.ModuleList(decoders)
# 1x1x1 conv + simple ReLU in the final convolution
self.final_conv = SingleConv(f_maps[0],
out_channels,
kernel_size=1,
order='cr',
padding=0)
def __init__(self,
in_channels,
out_channels,
final_sigmoid,
basic_module,
f_maps=64,
layer_order='gcr',
num_groups=8,
num_levels=4,
is_segmentation=True,
testing=False,
conv_kernel_size=3,
pool_kernel_size=2,
conv_padding=1,
**kwargs):
super(Abstract3DUNet, self).__init__()
self.testing = testing
if isinstance(f_maps, int):
f_maps = number_of_features_per_level(f_maps,
num_levels=num_levels)
# create encoder path consisting of Encoder modules. Depth of the encoder is equal to `len(f_maps)`
encoders = []
for i, out_feature_num in enumerate(f_maps):
if i == 0:
encoder = Encoder(
in_channels,
out_feature_num,
apply_pooling=False, # skip pooling in the firs encoder
basic_module=basic_module,
conv_layer_order=layer_order,
conv_kernel_size=conv_kernel_size,
num_groups=num_groups,
padding=conv_padding)
else:
# TODO: adapt for anisotropy in the data, i.e. use proper pooling kernel to make the data isotropic after 1-2 pooling operations
encoder = Encoder(f_maps[i - 1],
out_feature_num,
basic_module=basic_module,
conv_layer_order=layer_order,
conv_kernel_size=conv_kernel_size,
num_groups=num_groups,
pool_kernel_size=pool_kernel_size,
padding=conv_padding)
encoders.append(encoder)
self.encoders = nn.ModuleList(encoders)
# create decoder path consisting of the Decoder modules. The length of the decoder is equal to `len(f_maps) - 1`
decoders = []
reversed_f_maps = list(reversed(f_maps))
for i in range(len(reversed_f_maps) - 1):
if basic_module == DoubleConv:
in_feature_num = reversed_f_maps[i] + reversed_f_maps[i + 1]
else:
in_feature_num = reversed_f_maps[i]
out_feature_num = reversed_f_maps[i + 1]
# TODO: if non-standard pooling was used, make sure to use correct striding for transpose conv
# currently strides with a constant stride: (2, 2, 2)
decoder = Decoder(in_feature_num,
out_feature_num,
basic_module=basic_module,
conv_layer_order=layer_order,
conv_kernel_size=conv_kernel_size,
num_groups=num_groups,
padding=conv_padding)
decoders.append(decoder)
self.decoders = nn.ModuleList(decoders)
# in the last layer a 1×1 convolution reduces the number of output
# channels to the number of labels
self.final_conv = nn.Conv3d(f_maps[0], out_channels, 1)
if is_segmentation:
# semantic segmentation problem
if final_sigmoid:
self.final_activation = nn.Sigmoid()
else:
self.final_activation = nn.Softmax(dim=1)
else:
# regression problem
self.final_activation = None
def __init__(self,
in_channels,
out_channels,
final_sigmoid,
f_maps=32,
conv_layer_order='cge',
num_groups=8,
skip_final_activation=False,
**kwargs):
super(ResidualUNet3D, self).__init__()
if isinstance(f_maps, int):
# use 5 levels in the encoder path as suggested in the paper
f_maps = create_feature_maps(f_maps, number_of_fmaps=5)
# create encoder path consisting of Encoder modules. The length of the encoder is equal to `len(f_maps)`
# uses ExtResNetBlock as a basic_module for the Encoder
encoders = []
for i, out_feature_num in enumerate(f_maps):
if i == 0:
encoder = Encoder(in_channels,
out_feature_num,
apply_pooling=False,
basic_module=ExtResNetBlock,
conv_layer_order=conv_layer_order,
num_groups=num_groups)
else:
encoder = Encoder(f_maps[i - 1],
out_feature_num,
basic_module=ExtResNetBlock,
conv_layer_order=conv_layer_order,
num_groups=num_groups)
encoders.append(encoder)
self.encoders = nn.ModuleList(encoders)
# create decoder path consisting of the Decoder modules. The length of the decoder is equal to `len(f_maps) - 1`
# uses ExtResNetBlock as a basic_module for the Decoder
decoders = []
reversed_f_maps = list(reversed(f_maps))
for i in range(len(reversed_f_maps) - 1):
decoder = Decoder(reversed_f_maps[i],
reversed_f_maps[i + 1],
basic_module=ExtResNetBlock,
conv_layer_order=conv_layer_order,
num_groups=num_groups)
decoders.append(decoder)
self.decoders = nn.ModuleList(decoders)
# in the last layer a 1×1 convolution reduces the number of output
# channels to the number of labels
self.final_conv = nn.Conv3d(f_maps[0], out_channels, 1)
if not skip_final_activation:
if final_sigmoid:
self.final_activation = nn.Sigmoid()
else:
self.final_activation = nn.Softmax(dim=1)
else:
self.final_activation = None
