biapy.models.attention_unet

Attention U-Net Implementation.

This module implements a 2D/3D U-Net architecture with attention gates for improved feature learning in medical image segmentation and other computer vision tasks.

The implementation is based on the paper: “Attention U-Net: Learning Where to Look for the Pancreas” https://arxiv.org/abs/1804.03999

The attention mechanism helps the model focus on relevant spatial regions while suppressing irrelevant background activations, leading to improved segmentation performance especially in medical imaging applications.

Classes:

Attention_U_Net: Main U-Net architecture with attention gates

class biapy.models.attention_unet.Attention_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', z_down=[2, 2, 2, 2], yx_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 Attention 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