biapy.models.seunet

This module implements the Squeeze-and-Excitation U-Net (SE-U-Net) architecture,a variant of the classic U-Net enhanced with Squeeze-and-Excitation (SE) blocks.

The SE-U-Net is designed for various image analysis tasks, particularly dense prediction problems like image segmentation, in both 2D and 3D. It integrates SE blocks into its convolutional layers to enable the network to perform dynamic channel-wise feature recalibration, thereby improving the quality of learned representations.

Key components and functionalities include:

Classes:

  • SE_U_Net: The main Squeeze-and-Excitation U-Net model, comprising an encoder (downsampling path), a decoder (upsampling path), and skip connections, all enhanced with SE blocks.

This module leverages several building blocks defined in biapy.models.blocks, such as DoubleConvBlock, UpBlock, ConvBlock, ProjectionHead, and utility functions for normalization (get_norm_2d, get_norm_3d).

Reference: Squeeze and Excitation Networks.

Image representation:

../../_images/unet.png

Image created with PlotNeuralNet.

class biapy.models.seunet.SE_U_Net(image_shape=(256, 256, 1), activation='ELU', feature_maps=[32, 64, 128, 256], drop_values=[0.1, 0.1, 0.1, 0.1], normalization='none', k_size=3, upsample_layer='convtranspose', yx_down=[2, 2, 2, 2], z_down=[2, 2, 2, 2], output_channels=[1], separated_decoders=False, output_channel_info=['F'], explicit_activations: bool = False, head_activations: List[str] = ['ce_sigmoid'], upsampling_factor=(), upsampling_position='pre', isotropy=False, larger_io=True, contrast: bool = False, contrast_proj_dim: int = 256, return_one_tensor: bool = False)[source]

Bases: Module

Create 2D/3D U-Net with Squeeze-and-Excitation (SE) blocks.

This model extends the classic U-Net architecture by incorporating Squeeze-and-Excitation (SE) modules within its convolutional blocks. This design aims to improve feature learning and propagation by allowing the network to perform dynamic channel-wise feature recalibration, leading to better performance in dense prediction tasks like image segmentation.

Reference: Squeeze and Excitation Networks.

forward(x) → Dict | Tensor[source]

Forward pass of the model.

Parameters:

x (torch.Tensor) – Input tensor of shape (batch_size, channels, height, width) for 2D or (batch_size, channels, depth, height, width) for 3D.

Returns:

Model output. Returns a dictionary if multi-head or contrastive outputs are enabled, otherwise returns the main prediction tensor.

Return type:

Dict or torch.Tensor