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=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, 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