Unet
- class colibri.models.unet.Unet(in_channels=1, out_channels=1, features=[32, 64, 128, 256], last_activation='sigmoid', **kwargs)[source]
Bases:
Module
Unet Model
The U-Net model is a fully convolutional neural network initially proposed for biomedical image segmentation. Similar to the autoencoder model, the U-Net model consists of an encoder and a decoder. The encoder extracts features from the input image, while the decoder upsamples these features to the original image size. During the upsampling process, the decoder uses skip connections to concatenate features from the encoder with the upsampled features. These skip connections help preserve the spatial information of the input image, which is the key difference between the U-Net and the autoencoder model.
Implementation based on the formulation of authors in https://doi.org/10.1007/978-3-319-24574-4_28
- Parameters:
in_channels (int) – number of input channels
out_channels (int) – number of output channels
features (list, optional) – number of features in each level of the Unet. Defaults to [32, 64, 128, 256].
last_activation (str, optional) – activation function for the last layer. Defaults to ‘sigmoid’.
- Returns:
Unet model
- Return type:
torch.nn.Module