biapy.models

This package (biapy.models) is responsible for building and managing deep learning models within the BiaPy framework.

It provides functionalities to:

Dynamically build models: Select and instantiate various neural network architectures (e.g., U-Net, ResUNet, ViT, ConvNeXt variants, etc.) based on configuration settings.
Integrate with BioImage Model Zoo (BMZ): Facilitate the loading and compatibility checking of pre-trained models from the BioImage Model Zoo, enabling easy reuse of community-contributed models.
Extract model source code: Collect the necessary source code for a given model and its dependencies, which is crucial for reproducibility and export functionalities.

The module handles different problem types (e.g., semantic segmentation, super-resolution, classification) and adapts model configurations (e.g., 2D/3D, input/output channels, normalization, dropout) accordingly.

biapy.models.build_model(cfg: CfgNode, output_channels: List[int], output_channel_info: List[str], head_activations: List[str], device: device) → Tuple[Module, str, Dict, set, List[str], Dict, Tuple[int, ...]][source]

Build selected model.

Parameters:

cfg (YACS CN object) – Configuration.
output_channels (List[int]) – Number of output channels for each head.
output_channel_info (List[str]) – Information about each output channel.
head_activations (List[str]) – Activation functions for each output head.
device (Torch device) – Using device. Most commonly “cpu” or “cuda” for GPU, but also potentially “mps”, “xpu”, “xla” or “meta”.

Returns:

model – Selected model.

Return type:

Pytorch model

biapy.models.init_embedding_output(model: Module, n_sigma: int = 2)[source]

Initialize the output layer of the model for embedding.

Parameters:

model (nn.Module) – The model whose output layer needs to be initialized.
n_sigma (int) – Number of sigma channels to initialize.

biapy.models.extract_model(dependency_queue: deque, model_file: str) → Tuple[Dict[str, str], set, List[str]][source]

Extract the source code of the model and its dependencies, ensuring dependencies are ordered before the definition that uses them.

Parameters:

dependency_queue (deque) – Queue of model dependencies to be processed.
model_file (str) – Path to the main model file.

Returns:

collected_sources (dict) – Dictionary containing the source code of the collected model dependencies, ordered so that dependencies appear before the main model.
all_import_lines (set) – Set of all external import lines found in the model and its dependencies.
scanned_files (list) – List of all files that were scanned for dependencies.

biapy.models.merge_import_lines(import_lines: List[str]) → List[str][source]

Merge import lines by grouping them by module and sorting names within each module.

Parameters:: import_lines (list of str) – List of import lines to be merged.
Returns:: merged – Merged import lines, sorted and grouped by module.
Return type:: list of str

biapy.models.adapt_bmz_model_kwargs(model_kwargs: Dict, model_to_consume: bool) → Dict[source]

Adapt BMZ model arguments to be compatible with BiaPy’s model building functions.

Parameters:

model_kwargs (dict) – Dictionary of model arguments to be adapted.
model_to_consume (bool) – Whether the model is being adapted for consumption (True) or for exporting (False).

Returns:

adapted_args – Dictionary of adapted arguments ready to be passed to the model building function.

Return type:

dict

biapy.models.get_bmz_model_kwargs(model: ModelDescr | ModelDescr) → Dict[source]

Get the PyTorch state dict weight specification from a BMZ model description.

Parameters:: model (ModelDescr_v0_4 | ModelDescr_v0_5) – BMZ model description.
Returns:: model_kwargs – Dictionary of model arguments extracted from the BMZ model description, ready to be adapted for BiaPy’s model building functions.
Return type:: dict

biapy.models.update_bmz_model_kwargs_to_biapy(new_biapy_model_kwargs: Dict, model: ModelDescr | ModelDescr)[source]

Build a model from Bioimage Model Zoo (BMZ).

Parameters:

new_biapy_model_kwargs (dict) – Dictionary of model arguments as expected by BiaPy’s model building functions.
model (ModelDescr) – BMZ model RDF that contains all the information of the model.

Returns:

bmz_model_kwargs – Updated dictionary of model arguments to be used in the BMZ model RDF, ensuring compatibility with BiaPy.

Return type:

dict

biapy.models.build_bmz_model(cfg: CfgNode, model: ModelDescr | ModelDescr, device: device) → Module[source]

Build a model from Bioimage Model Zoo (BMZ).

Parameters:

cfg (YACS configuration) – Running configuration.
model (ModelDescr) – BMZ model RDF that contains all the information of the model.
device (Torch device) – Device used.

Returns:

model_instance – Torch model.

Return type:

Torch model

biapy.models.is_biapy_model(model: ModelDescr | ModelDescr) → bool[source]

Check if a model is a BiaPy model by looking for:

the presence of “danifranco” in the GitHub username or “Daniel Franco” in the author name.
the presence of “biapy” in the model tags.
the presence of a citation with the text “BiaPy: accessible deep learning on bioimages” in the citations of the model.

Parameters:: model (ModelDescr_v0_4 | ModelDescr_v0_5) – The model to check.
Returns:: True if the model is a BiaPy model, False otherwise.
Return type:: bool

biapy.models.find_bmz_models(model_ID: str | None = None, url: str = 'https://hypha.aicell.io/bioimage-io/artifacts/bioimage.io/children?limit=1000000', timeout: int = 30)[source]

Query the BioImage.IO Hypha API for models and return those whose nickname/id/rdf_source contains model_ID (case-insensitive).

Returns list of dicts with: id, alias, nickname, rdf_source, version, format_version, artifact_path (id used as path), and a few handy urls.

Parameters:

model_ID (str) – Model identifier. It can be either its DOI or nickname. Leave it as None to get all available models.
url (str) – URL to the BioImage.IO Hypha API endpoint to query for models.
timeout (int) – Timeout for the HTTP request in seconds.

Returns:

out – List of dictionaries containing model information. Each dictionary has the following keys: id, alias, nickname, rdf_source, version, format_version, artifact_path, urls (which contains covers and documentation URLs), and raw (the original item from the API response).

Return type:

list of dict

biapy.models.check_bmz_args(model_ID: str, cfg: CfgNode) → Tuple[List, Dict][source]

Check user’s provided BMZ arguments.

Parameters:

model_ID (str) – Model identifier. It can be either its DOI or nickname.
cfg (YACS configuration) – Running configuration.

Returns:

preproc_info – Preprocessing names that the model is using.

Return type:

dict

biapy.models.check_bmz_model_compatibility(model_rdf: Dict, workflow_specs: Dict | None = None) → Tuple[List, bool, str, Dict][source]

Check one model compatibility with BiaPy by looking at its RDF file provided by BMZ. This function is the one used in BMZ’s continuous integration with BiaPy.

Parameters:

model_rdf (dict) – BMZ model RDF that contains all the information of the model.
workflow_specs (dict) – Specifications of the workflow. If not provided all possible models will be considered.

Returns:

preproc_info (dict) – Preprocessing names that the model is using.
error (bool) – Whether it there is a problem to consume the model in BiaPy or not.
reason_message (str) – Reason why the model can not be consumed if there is any.

biapy.models.build_torchvision_model(cfg: CfgNode, device: device) → Tuple[Module, Callable][source]

Build and adapt a model from the torchvision.models library based on the configuration.

This function dynamically loads a pre-trained model from torchvision.models (e.g., ResNet, DeepLabV3, MaskRCNN, etc.) as specified in the configuration. It then adapts the final output layer(s) of the model to match the number of classes or output channels required by the specific problem type (e.g., classification, semantic segmentation, instance segmentation).

Parameters:

cfg (YACS CN object) – The configuration object. Key parameters used are:
- cfg.MODEL.TORCHVISION_MODEL_NAME: Name of the torchvision model to load (e.g., “resnet50”, “deeplabv3_resnet101”, “maskrcnn_resnet50_fpn”, “quantized_resnet50”).
- cfg.PROBLEM.TYPE: Type of problem (e.g., “CLASSIFICATION”, “SEMANTIC_SEG”, “INSTANCE_SEG”, “DETECTION”) to determine model adaptation logic.
- cfg.DATA.N_CLASSES: Number of output classes required for the problem.
- cfg.DATA.PATCH_SIZE: Input patch size, used for generating the model summary.
- cfg.PROBLEM.NDIM: Number of input dimensions (“2D” or “3D”).
device (torch.device) – The PyTorch device (e.g., “cpu”, “cuda”, “mps”) on which the model will be loaded and run.

Returns:

model (nn.Module) – The instantiated and adapted PyTorch model from torchvision.
transforms (Callable) – A callable representing the default preprocessing transforms associated with the loaded torchvision model’s weights. This should be applied to input images before feeding them to the model.

Notes

Models are loaded with their DEFAULT pre-trained weights from torchvision.
The final layer adaptation logic is specific to common torchvision model structures for classification, semantic segmentation, and instance segmentation.
For classification, the final linear layer is replaced. A warning is printed if the number of classes differs from ImageNet’s default (1000).
For semantic segmentation, the final convolutional layer(s) of the classifier and auxiliary classifier (if present) are replaced. A warning is printed if the number of classes differs from Pascal VOC’s default (21).
For instance segmentation (MaskRCNN), the box predictor’s classification score head and the mask predictor’s final convolutional layer are replaced. A warning is printed if the number of classes differs from COCO’s default (91).
Special handling is included for squeezenet and lraspp_mobilenet_v3_large due to their unique head structures.
For maxvit_t in classification, a fixed sample input size of (1, 3, 224, 224) is used for the model summary.
This function assumes the necessary torchvision models and their default weights are installed and accessible.

biapy.models.can_import_env_deps(yaml_reader, import_overrides={'pyyaml': 'yaml', 'scikit-learn': 'sklearn'}, allowlist: Iterable[str] | None = {'pytorch', 'pytorch-cuda', 'pytorch-mutex', 'torch', 'torchvision'}) → Tuple[bool, str][source]

Check if all dependencies listed in a conda-style environment yaml file can be imported. Dependencies whose distribution name is in allowlist are ignored.

Parameters:

yaml_reader (file-like object) – Provides the content of a conda-style environment yaml file.
import_overrides (dict, optional) – Map dist name -> import name (e.g. {‘pyyaml’: ‘yaml’}).
allowlist (Iterable[str], optional) – Dist names to ignore if they fail import/version checks (case-insensitive). Example: {“pytorch”, “torch”, “pytorch-cuda”}

Returns:

ok (bool)
msg (str)

biapy.models.get_last_layer_info(model: Module) → Dict[str, Any][source]

Recursively finds the last layer of a model and checks if it’s an activation.

Parameters:: model (nn.Module) – The PyTorch model to analyze.
Returns:: A dictionary containing: - “layer_object”: The last layer object found in the model. - “layer_type”: The type name of the last layer. - “is_activation”: A boolean indicating whether the last layer is a common activation function.
Return type:: dict

Submodules

Module	Description
`biapy.models.attention_unet`	U-Net variant with attention gates for improved feature propagation.
`biapy.models.dfcan`	Deep Fourier Channel Attention Network for image restoration.
`biapy.models.edsr`	Enhanced Deep Super-Resolution (EDSR) model.
`biapy.models.efficientnet`	EfficientNet backbone for image classification and feature extraction.
`biapy.models.hrnet`	High-Resolution Network (HRNet) for semantic segmentation.
`biapy.models.mae`	Masked Autoencoder (MAE) for self-supervised pretraining.
`biapy.models.multiresunet`	MultiResUNet for multi-scale feature fusion in segmentation.
`biapy.models.rcan`	Residual Channel Attention Network for image super-resolution.
`biapy.models.resunet`	Residual U-Net for segmentation tasks.
`biapy.models.resunet_plus_plus`	ResUNet++: residual U-Net variant with attention and dense blocks.
`biapy.models.resunet_se`	Residual U-Net with Squeeze-and-Excitation blocks.
`biapy.models.seunet`	U-Net variant incorporating Squeeze-and-Excitation blocks.
`biapy.models.simple_cnn`	Simple convolutional neural network for baseline tasks.
`biapy.models.unet`	Classic U-Net architecture for image segmentation.
`biapy.models.unetr`	U-Net Transformer (UNETR) architecture.
`biapy.models.unext_v1`	UNeXt version 1: lightweight U-Net with MLP-like blocks.
`biapy.models.unext_v2`	UNeXt version 2: refined lightweight U-Net variant.
`biapy.models.vit`	Vision Transformer (ViT) backbone for image analysis.
`biapy.models.wdsr`	Wide Activation Super-Resolution (WDSR) model.
`biapy.models.blocks`	Building blocks and layers used across multiple models.
`biapy.models.bmz_utils`	Utility functions for BMZ-based operations.
`biapy.models.memory_bank`	Memory bank module for contrastive learning.
`biapy.models.tr_layers`	Transformer layers and utilities.