careamics 0.0.19__py3-none-any.whl

careamics/model_io/bmz_io.py

@@ -0,0 +1,251 @@
+"""Function to export to the BioImage Model Zoo format."""
+
+import tempfile
+from pathlib import Path
+from typing import Union
+
+import numpy as np
+from bioimageio.core import load_model_description, test_model
+from bioimageio.spec import ValidationSummary, save_bioimageio_package
+from pydantic import HttpUrl
+from torch import __version__ as PYTORCH_VERSION
+from torch import load, save
+from torchvision import __version__ as TORCHVISION_VERSION
+
+from careamics.config import Configuration, load_configuration, save_configuration
+from careamics.config.support import SupportedArchitecture
+from careamics.lightning.lightning_module import FCNModule, VAEModule
+from careamics.utils.version import get_careamics_version
+
+from .bioimage import (
+    create_env_text,
+    create_model_description,
+    extract_model_path,
+)
+from .bioimage.cover_factory import create_cover
+
+
+def _export_state_dict(
+    model: Union[FCNModule, VAEModule], path: Union[Path, str]
+) -> Path:
+    """
+    Export the model state dictionary to a file.
+
+    Parameters
+    ----------
+    model : CAREamicsKiln
+        CAREamics model to export.
+    path : Union[Path, str]
+        Path to the file where to save the model state dictionary.
+
+    Returns
+    -------
+    Path
+        Path to the saved model state dictionary.
+    """
+    path = Path(path)
+
+    # make sure it has the correct suffix
+    if path.suffix not in ".pth":
+        path = path.with_suffix(".pth")
+
+    # save model state dictionary
+    # we save through the torch model itself to avoid the initial "model." in the
+    # layers naming, which is incompatible with the way the BMZ load torch state dicts
+    save(model.model.state_dict(), path)
+
+    return path
+
+
+def _load_state_dict(
+    model: Union[FCNModule, VAEModule], path: Union[Path, str]
+) -> None:
+    """
+    Load a model from a state dictionary.
+
+    Parameters
+    ----------
+    model : CAREamicsKiln
+        CAREamics model to be updated with the weights.
+    path : Union[Path, str]
+        Path to the model state dictionary.
+    """
+    path = Path(path)
+
+    # load model state dictionary
+    # same as in _export_state_dict, we load through the torch model to be compatible
+    # witht bioimageio.core expectations for a torch state dict
+    state_dict = load(path)
+    model.model.load_state_dict(state_dict)
+
+
+# TODO break down in subfunctions
+def export_to_bmz(
+    model: Union[FCNModule, VAEModule],
+    config: Configuration,
+    path_to_archive: Union[Path, str],
+    model_name: str,
+    general_description: str,
+    data_description: str,
+    authors: list[dict],
+    input_array: np.ndarray,
+    output_array: np.ndarray,
+    covers: list[Union[Path, str]] | None = None,
+    channel_names: list[str] | None = None,
+    model_version: str = "0.1.0",
+) -> None:
+    """Export the model to BioImage Model Zoo format.
+
+    Arrays are expected to be SC(Z)YX with singleton dimensions allowed for S and C.
+
+    `model_name` should consist of letters, numbers, dashes, underscores and parentheses
+    only.
+
+    Parameters
+    ----------
+    model : CAREamicsModule
+        CAREamics model to export.
+    config : Configuration
+        Model configuration.
+    path_to_archive : Union[Path, str]
+        Path to the output file.
+    model_name : str
+        Model name.
+    general_description : str
+        General description of the model.
+    data_description : str
+        Description of the data the model was trained on.
+    authors : list[dict]
+        Authors of the model.
+    input_array : np.ndarray
+        Input array, should not have been normalized.
+    output_array : np.ndarray
+        Output array, should have been denormalized.
+    covers : list of pathlib.Path or str, default=None
+        Paths to the cover images.
+    channel_names : Optional[list[str]], optional
+        Channel names, by default None.
+    model_version : str, default="0.1.0"
+        Model version.
+    """
+    path_to_archive = Path(path_to_archive)
+
+    if path_to_archive.suffix != ".zip":
+        raise ValueError(
+            f"Path to archive must point to a zip file, got {path_to_archive}."
+        )
+
+    if not path_to_archive.parent.exists():
+        path_to_archive.parent.mkdir(parents=True, exist_ok=True)
+
+    # versions
+    careamics_version = get_careamics_version()
+
+    # save files in temporary folder
+    with tempfile.TemporaryDirectory() as tmpdirname:
+        temp_path = Path(tmpdirname)
+
+        # create environment file
+        # TODO move in bioimage module
+        env_path = temp_path / "environment.yml"
+        env_path.write_text(create_env_text(PYTORCH_VERSION, TORCHVISION_VERSION))
+
+        # export input and ouputs
+        inputs = temp_path / "inputs.npy"
+        np.save(inputs, input_array)
+        outputs = temp_path / "outputs.npy"
+        np.save(outputs, output_array)
+
+        # export configuration
+        config_path = save_configuration(config, temp_path / "careamics.yaml")
+
+        # export model state dictionary
+        weight_path = _export_state_dict(model, temp_path / "weights.pth")
+
+        # export cover if necesary
+        if covers is None:
+            covers = [create_cover(temp_path, input_array, output_array)]
+
+        # create model description
+        model_description = create_model_description(
+            config=config,
+            name=model_name,
+            general_description=general_description,
+            data_description=data_description,
+            authors=authors,
+            inputs=inputs,
+            outputs=outputs,
+            weights_path=weight_path,
+            torch_version=PYTORCH_VERSION,
+            careamics_version=careamics_version,
+            config_path=config_path,
+            env_path=env_path,
+            covers=covers,
+            channel_names=channel_names,
+            model_version=model_version,
+        )
+
+        # test model description
+        test_kwargs = {}
+        if hasattr(model_description, "config") and isinstance(
+            model_description.config, dict
+        ):
+            bioimageio_config = model_description.config.get("bioimageio", {})
+            test_kwargs = bioimageio_config.get("test_kwargs", {}).get(
+                "pytorch_state_dict", {}
+            )
+
+        summary: ValidationSummary = test_model(model_description, **test_kwargs)
+        if summary.status == "failed":
+            raise ValueError(f"Model description test failed: {summary}")
+
+        # save bmz model
+        save_bioimageio_package(model_description, output_path=path_to_archive)
+
+
+def load_from_bmz(
+    path: Union[Path, str, HttpUrl],
+) -> tuple[Union[FCNModule, VAEModule], Configuration]:
+    """Load a model from a BioImage Model Zoo archive.
+
+    Parameters
+    ----------
+    path : Path, str or HttpUrl
+        Path to the BioImage Model Zoo archive. A Http URL must point to a downloadable
+        location.
+
+    Returns
+    -------
+    FCNModel or VAEModel
+        The loaded CAREamics model.
+    Configuration
+        The loaded CAREamics configuration.
+
+    Raises
+    ------
+    ValueError
+        If the path is not a zip file.
+    """
+    # load description, this creates an unzipped folder next to the archive
+    model_desc = load_model_description(path)
+
+    # extract paths
+    weights_path, config_path = extract_model_path(model_desc)
+
+    # load configuration
+    config = load_configuration(config_path)
+
+    # create careamics lightning module
+    if config.algorithm_config.model.architecture == SupportedArchitecture.UNET:
+        model = FCNModule(algorithm_config=config.algorithm_config)
+    elif config.algorithm_config.model.architecture == SupportedArchitecture.LVAE:
+        model = VAEModule(algorithm_config=config.algorithm_config)
+    else:
+        raise ValueError(
+            f"Unsupported architecture {config.algorithm_config.model.architecture}"
+        )  # TODO ugly ?
+
+    # load model state dictionary
+    _load_state_dict(model, weights_path)
+
+    return model, config

careamics/model_io/model_io_utils.py

@@ -0,0 +1,95 @@
+"""Utility functions to load pretrained models."""
+
+from pathlib import Path
+from typing import Union
+
+import torch
+
+from careamics.config import Configuration
+from careamics.lightning.lightning_module import FCNModule, VAEModule
+from careamics.model_io.bmz_io import load_from_bmz
+from careamics.utils import check_path_exists
+
+
+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}."
+        )
+
+
+def _load_checkpoint(
+    path: Union[Path, str],
+) -> tuple[Union[FCNModule, VAEModule], Configuration]:
+    """
+    Load a model from a checkpoint and return both model and configuration.
+
+    Parameters
+    ----------
+    path : Union[Path, str]
+        Path to the checkpoint.
+
+    Returns
+    -------
+    tuple[CAREamicsKiln, Configuration]
+        tuple of CAREamics model and its configuration.
+
+    Raises
+    ------
+    ValueError
+        If the checkpoint file does not contain hyper parameters (configuration).
+    """
+    # load checkpoint
+    # here we might run into issues between devices
+    # see https://pytorch.org/tutorials/recipes/recipes/save_load_across_devices.html
+    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+    checkpoint: dict = torch.load(path, map_location=device)
+
+    # attempt to load configuration
+    try:
+        cfg_dict = checkpoint["hyper_parameters"]
+    except KeyError as e:
+        raise ValueError(
+            f"Invalid checkpoint file. No `hyper_parameters` found in the "
+            f"checkpoint: {checkpoint.keys()}"
+        ) from e
+
+    if cfg_dict["algorithm_config"]["model"]["architecture"] == "UNet":
+        model = FCNModule.load_from_checkpoint(path)
+    elif cfg_dict["algorithm_config"]["model"]["architecture"] == "LVAE":
+        model = VAEModule.load_from_checkpoint(path)
+    else:
+        raise ValueError(
+            "Invalid model architecture: "
+            f"{cfg_dict['algorithm_config']['model']['architecture']}"
+        )
+
+    return model, Configuration(**cfg_dict)

careamics/models/__init__.py

@@ -0,0 +1,5 @@
+"""Models package."""
+
+__all__ = ["model_factory"]
+
+from .model_factory import model_factory

careamics/models/activation.py

@@ -0,0 +1,40 @@
+"""Activations for CAREamics models."""
+
+from collections.abc import Callable
+from typing import Union
+
+import torch.nn as nn
+
+from ..config.support import SupportedActivation
+
+
+def get_activation(activation: Union[SupportedActivation, str]) -> Callable:
+    """
+    Get activation function.
+
+    Parameters
+    ----------
+    activation : str
+        Activation function name.
+
+    Returns
+    -------
+    Callable
+        Activation function.
+    """
+    if activation == SupportedActivation.RELU:
+        return nn.ReLU()
+    elif activation == SupportedActivation.ELU:
+        return nn.ELU()
+    elif activation == SupportedActivation.LEAKYRELU:
+        return nn.LeakyReLU()
+    elif activation == SupportedActivation.TANH:
+        return nn.Tanh()
+    elif activation == SupportedActivation.SIGMOID:
+        return nn.Sigmoid()
+    elif activation == SupportedActivation.SOFTMAX:
+        return nn.Softmax(dim=1)
+    elif activation == SupportedActivation.NONE:
+        return nn.Identity()
+    else:
+        raise ValueError(f"Activation {activation} not supported.")