biapy.models.resunet

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

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

Reference:

Road Extraction by Deep Residual U-Net https://ieeexplore.ieee.org/abstract/document/8309343

class biapy.models.resunet.ResUNet(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 Residual 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