biapy.models.unet

U-Net model definition for 2D/3D biomedical image tasks.

This module provides the U_Net class, which implements a configurable U-Net architecture with support for 2D and 3D data, super-resolution, multi-head outputs, and optional contrastive learning head.

Reference:

U-Net: Convolutional Networks for Biomedical Image Segmentation https://link.springer.com/chapter/10.1007/978-3-319-24574-4_28

class biapy.models.unet.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)

Bases: Module

Configurable 2D/3D U-Net model for image segmentation and super-resolution.

Supports multi-head outputs and optional contrastive learning head.

forward(x) → Dict | Tensor

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