Skip to content

model_io_utils

Utility functions to load pretrained models.

load_pretrained(path) #

Load a pretrained model from a checkpoint or a BioImage Model Zoo model.

Expected formats are .ckpt or .zip files.

Parameters:

Name Type Description Default
path Union[Path, str]

Path to the pretrained model.

required

Returns:

Type Description
Tuple[CAREamicsKiln, Configuration]

Tuple of CAREamics model and its configuration.

Raises:

Type Description
ValueError

If the model format is not supported.

Source code in src/careamics/model_io/model_io_utils.py
def load_pretrained(
    path: Union[Path, str]
) -> Tuple[Union[FCNModule, VAEModule], Configuration]:
    """
    Load a pretrained model from a checkpoint or a BioImage Model Zoo model.

    Expected formats are .ckpt or .zip files.

    Parameters
    ----------
    path : Union[Path, str]
        Path to the pretrained model.

    Returns
    -------
    Tuple[CAREamicsKiln, Configuration]
        Tuple of CAREamics model and its configuration.

    Raises
    ------
    ValueError
        If the model format is not supported.
    """
    path = check_path_exists(path)

    if path.suffix == ".ckpt":
        return _load_checkpoint(path)
    elif path.suffix == ".zip":
        return load_from_bmz(path)
    else:
        raise ValueError(
            f"Invalid model format. Expected .ckpt or .zip, got {path.suffix}."
        )