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

forward(x)[source]

Args :

x (torch.Tensor): Input tensor

Returns:

Output tensor

Return type:

torch.Tensor