careamics 0.1.0rc2__py3-none-any.whl → 0.1.0rc3__py3-none-any.whl

careamics/model_io/bioimage/bioimage_utils.py

@@ -0,0 +1,48 @@
+"""Bioimage.io utils."""
+from pathlib import Path
+from typing import Union
+
+
+def get_unzip_path(zip_path: Union[Path, str]) -> Path:
+    """Generate unzipped folder path from the bioimage.io model path.
+
+    Parameters
+    ----------
+    zip_path : Path
+        Path to the bioimage.io model.
+
+    Returns
+    -------
+    Path
+        Path to the unzipped folder.
+    """
+    zip_path = Path(zip_path)
+
+    return zip_path.parent / (str(zip_path.name) + ".unzip")
+
+
+def create_env_text(pytorch_version: str) -> str:
+    """Create environment text for the bioimage model.
+
+    Parameters
+    ----------
+    pytorch_version : str
+        Pytorch version.
+
+    Returns
+    -------
+    str
+        Environment text.
+    """
+    env = (
+        f"name: careamics\n"
+        f"dependencies:\n"
+        f"  - python=3.8\n"
+        f"  - pytorch={pytorch_version}\n"
+        f"  - torchvision={pytorch_version}\n"
+        f"  - pip\n"
+        f"  - pip:\n"
+        f"    - git+https://github.com/CAREamics/careamics.git@dl4mia\n"
+    )
+    # TODO from pip with package version
+    return env

careamics/model_io/bioimage/model_description.py

@@ -0,0 +1,318 @@
+"""Module use to build BMZ model description."""
+from pathlib import Path
+from typing import List, Optional, Tuple, Union
+
+import numpy as np
+from bioimageio.spec.model.v0_5 import (
+    ArchitectureFromLibraryDescr,
+    Author,
+    AxisBase,
+    AxisId,
+    BatchAxis,
+    ChannelAxis,
+    EnvironmentFileDescr,
+    FileDescr,
+    FixedZeroMeanUnitVarianceDescr,
+    FixedZeroMeanUnitVarianceKwargs,
+    Identifier,
+    InputTensorDescr,
+    ModelDescr,
+    OutputTensorDescr,
+    PytorchStateDictWeightsDescr,
+    SpaceInputAxis,
+    SpaceOutputAxis,
+    TensorId,
+    Version,
+    WeightsDescr,
+)
+
+from careamics.config import Configuration, DataModel
+
+from ._readme_factory import readme_factory
+
+
+def _create_axes(
+    array: np.ndarray,
+    data_config: DataModel,
+    channel_names: Optional[List[str]] = None,
+    is_input: bool = True,
+) -> List[AxisBase]:
+    """Create axes description.
+
+    Array shape is expected to be SC(Z)YX.
+
+    Parameters
+    ----------
+    array : np.ndarray
+        Array.
+    data_config : DataModel
+        CAREamics data configuration.
+    channel_names : Optional[List[str]], optional
+        Channel names, by default None.
+    is_input : bool, optional
+        Whether the axes are input axes, by default True.
+
+    Returns
+    -------
+    List[AxisBase]
+        List of axes description.
+
+    Raises
+    ------
+    ValueError
+        If channel names are not provided when channel axis is present.
+    """
+    # axes have to be SC(Z)YX
+    spatial_axes = data_config.axes.replace("S", "").replace("C", "")
+
+    # batch is always present
+    axes_model = [BatchAxis()]
+
+    if "C" in data_config.axes:
+        if channel_names is not None:
+            axes_model.append(
+                ChannelAxis(channel_names=[Identifier(name) for name in channel_names])
+            )
+        else:
+            raise ValueError(
+                f"Channel names must be provided if channel axis is present, axes: "
+                f"{data_config.axes}."
+            )
+    else:
+        # singleton channel
+        axes_model.append(ChannelAxis(channel_names=[Identifier("channel")]))
+
+    # spatial axes
+    for ind, axes in enumerate(spatial_axes):
+        if axes in ["X", "Y", "Z"]:
+            if is_input:
+                axes_model.append(
+                    SpaceInputAxis(id=AxisId(axes.lower()), size=array.shape[2 + ind])
+                )
+            else:
+                axes_model.append(
+                    SpaceOutputAxis(id=AxisId(axes.lower()), size=array.shape[2 + ind])
+                )
+
+    return axes_model
+
+
+def _create_inputs_ouputs(
+    input_array: np.ndarray,
+    output_array: np.ndarray,
+    data_config: DataModel,
+    input_path: Union[Path, str],
+    output_path: Union[Path, str],
+    channel_names: Optional[List[str]] = None,
+) -> Tuple[InputTensorDescr, OutputTensorDescr]:
+    """Create input and output tensor description.
+
+    Input and output paths must point to a `.npy` file.
+
+    Parameters
+    ----------
+    input_array : np.ndarray
+        Input array.
+    output_array : np.ndarray
+        Output array.
+    data_config : DataModel
+        CAREamics data configuration.
+    input_path : Union[Path, str]
+        Path to input .npy file.
+    output_path : Union[Path, str]
+        Path to output .npy file.
+    channel_names : Optional[List[str]], optional
+        Channel names, by default None.
+
+    Returns
+    -------
+    Tuple[InputTensorDescr, OutputTensorDescr]
+        Input and output tensor descriptions.
+    """
+    input_axes = _create_axes(input_array, data_config, channel_names)
+    output_axes = _create_axes(output_array, data_config, channel_names, False)
+
+    # mean and std
+    assert data_config.mean is not None, "Mean cannot be None."
+    assert data_config.std is not None, "Std cannot be None."
+    mean = data_config.mean
+    std = data_config.std
+
+    # and the mean and std required to invert the normalization
+    # CAREamics denormalization: x = y * (std + eps) + mean
+    # BMZ normalization : x = (y - mean') / (std' + eps)
+    # to apply the BMZ normalization as a denormalization step, we need:
+    eps = 1e-6
+    inv_mean = -mean / (std + eps)
+    inv_std = 1 / (std + eps) - eps
+
+    # create input/output descriptions
+    input_descr = InputTensorDescr(
+        id=TensorId("input"),
+        axes=input_axes,
+        test_tensor=FileDescr(source=input_path),
+        preprocessing=[
+            FixedZeroMeanUnitVarianceDescr(
+                kwargs=FixedZeroMeanUnitVarianceKwargs(mean=mean, std=std)
+            )
+        ],
+    )
+    output_descr = OutputTensorDescr(
+        id=TensorId("prediction"),
+        axes=output_axes,
+        test_tensor=FileDescr(source=output_path),
+        postprocessing=[
+            FixedZeroMeanUnitVarianceDescr(
+                kwargs=FixedZeroMeanUnitVarianceKwargs(  # invert normalization
+                    mean=inv_mean, std=inv_std
+                )
+            )
+        ],
+    )
+
+    return input_descr, output_descr
+
+
+def create_model_description(
+    config: Configuration,
+    name: str,
+    general_description: str,
+    authors: List[Author],
+    inputs: Union[Path, str],
+    outputs: Union[Path, str],
+    weights_path: Union[Path, str],
+    torch_version: str,
+    careamics_version: str,
+    config_path: Union[Path, str],
+    env_path: Union[Path, str],
+    channel_names: Optional[List[str]] = None,
+    data_description: Optional[str] = None,
+) -> ModelDescr:
+    """Create model description.
+
+    Parameters
+    ----------
+    config : Configuration
+        CAREamics configuration.
+    name : str
+        Name fo the model.
+    general_description : str
+        General description of the model.
+    authors : List[Author]
+        Authors of the model.
+    inputs : Union[Path, str]
+        Path to input .npy file.
+    outputs : Union[Path, str]
+        Path to output .npy file.
+    weights_path : Union[Path, str]
+        Path to model weights.
+    torch_version : str
+        Pytorch version.
+    careamics_version : str
+        CAREamics version.
+    config_path : Union[Path, str]
+        Path to model configuration.
+    env_path : Union[Path, str]
+        Path to environment file.
+    channel_names : Optional[List[str]], optional
+        Channel names, by default None.
+    data_description : Optional[str], optional
+        Description of the data, by default None.
+
+    Returns
+    -------
+    ModelDescr
+        Model description.
+    """
+    # documentation
+    doc = readme_factory(
+        config,
+        careamics_version=careamics_version,
+        data_description=data_description,
+    )
+
+    # inputs, outputs
+    input_descr, output_descr = _create_inputs_ouputs(
+        input_array=np.load(inputs),
+        output_array=np.load(outputs),
+        data_config=config.data_config,
+        input_path=inputs,
+        output_path=outputs,
+        channel_names=channel_names,
+    )
+
+    # weights description
+    architecture_descr = ArchitectureFromLibraryDescr(
+        import_from="careamics.models",
+        callable=f"{config.algorithm_config.model.architecture}",
+        kwargs=config.algorithm_config.model.model_dump(),
+    )
+
+    weights_descr = WeightsDescr(
+        pytorch_state_dict=PytorchStateDictWeightsDescr(
+            source=weights_path,
+            architecture=architecture_descr,
+            pytorch_version=Version(torch_version),
+            dependencies=EnvironmentFileDescr(source=env_path),
+        ),
+    )
+
+    # overall model description
+    model = ModelDescr(
+        name=name,
+        authors=authors,
+        description=general_description,
+        documentation=doc,
+        inputs=[input_descr],
+        outputs=[output_descr],
+        tags=config.get_algorithm_keywords(),
+        links=[
+            "https://github.com/CAREamics/careamics",
+            "https://careamics.github.io/latest/",
+        ],
+        license="BSD-3-Clause",
+        version="0.1.0",
+        weights=weights_descr,
+        attachments=[FileDescr(source=config_path)],
+        cite=config.get_algorithm_citations(),
+        config={  # conversion from float32 to float64 creates small differences...
+            "bioimageio": {
+                "test_kwargs": {
+                    "pytorch_state_dict": {
+                        "decimals": 2,  # ...so we relax the constraints on the decimals
+                    }
+                }
+            }
+        },
+    )
+
+    return model
+
+
+def extract_model_path(model_desc: ModelDescr) -> Tuple[Path, Path]:
+    """Return the relative path to the weights and configuration files.
+
+    Parameters
+    ----------
+    model_desc : ModelDescr
+        Model description.
+
+    Returns
+    -------
+    Tuple[Path, Path]
+        Weights and configuration paths.
+    """
+    weights_path = model_desc.weights.pytorch_state_dict.source.path
+
+    if len(model_desc.attachments) == 1:
+        config_path = model_desc.attachments[0].source.path
+    else:
+        for file in model_desc.attachments:
+            if file.source.path.suffix == ".yml":
+                config_path = file.source.path
+                break
+
+        if config_path is None:
+            raise ValueError("Configuration file not found.")
+
+    return weights_path, config_path

careamics/model_io/bmz_io.py

@@ -0,0 +1,231 @@
+"""Function to export to the BioImage Model Zoo format."""
+import tempfile
+from pathlib import Path
+from typing import List, Optional, Tuple, Union
+
+import numpy as np
+import pkg_resources
+from bioimageio.core import load_description, test_model
+from bioimageio.spec import ValidationSummary, save_bioimageio_package
+from torch import __version__, load, save
+
+from careamics.config import Configuration, load_configuration, save_configuration
+from careamics.config.support import SupportedArchitecture
+from careamics.lightning_module import CAREamicsKiln
+
+from .bioimage import (
+    create_env_text,
+    create_model_description,
+    extract_model_path,
+    get_unzip_path,
+)
+
+
+def _export_state_dict(model: CAREamicsKiln, 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: CAREamicsKiln, 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: CAREamicsKiln,
+    config: Configuration,
+    path: Union[Path, str],
+    name: str,
+    general_description: str,
+    authors: List[dict],
+    input_array: np.ndarray,
+    output_array: np.ndarray,
+    channel_names: Optional[List[str]] = None,
+    data_description: Optional[str] = None,
+) -> 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.
+
+    Parameters
+    ----------
+    model : CAREamicsKiln
+        CAREamics model to export.
+    config : Configuration
+        Model configuration.
+    path : Union[Path, str]
+        Path to the output file.
+    name : str
+        Model name.
+    general_description : str
+        General description of the model.
+    authors : List[dict]
+        Authors of the model.
+    input_array : np.ndarray
+        Input array.
+    output_array : np.ndarray
+        Output array.
+    channel_names : Optional[List[str]], optional
+        Channel names, by default None.
+    data_description : Optional[str], optional
+        Description of the data, by default None.
+
+    Raises
+    ------
+    ValueError
+        If the model is a Custom model.
+    """
+    path = Path(path)
+
+    # method is not compatible with Custom models
+    if config.algorithm_config.model.architecture == SupportedArchitecture.CUSTOM:
+        raise ValueError(
+            "Exporting Custom models to BioImage Model Zoo format is not supported."
+        )
+
+    # make sure that input and output arrays have the same shape
+    assert input_array.shape == output_array.shape, (
+        f"Input ({input_array.shape}) and output ({output_array.shape}) arrays "
+        f"have different shapes"
+    )
+
+    # make sure it has the correct suffix
+    if path.suffix not in ".zip":
+        path = path.with_suffix(".zip")
+
+    # versions
+    pytorch_version = __version__
+    careamics_version = pkg_resources.get_distribution("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))
+
+        # 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)
+
+        # export model state dictionary
+        weight_path = _export_state_dict(model, temp_path / "weights.pth")
+
+        # create model description
+        model_description = create_model_description(
+            config=config,
+            name=name,
+            general_description=general_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,
+            channel_names=channel_names,
+            data_description=data_description,
+        )
+
+        # test model description
+        summary: ValidationSummary = test_model(model_description)
+        if summary.status == "failed":
+            raise ValueError(f"Model description test failed: {summary}")
+
+        # save bmz model
+        save_bioimageio_package(model_description, output_path=path)
+
+
+def load_from_bmz(path: Union[Path, str]) -> Tuple[CAREamicsKiln, Configuration]:
+    """Load a model from a BioImage Model Zoo archive.
+
+    Parameters
+    ----------
+    path : Union[Path, str]
+        Path to the BioImage Model Zoo archive.
+
+    Returns
+    -------
+    Tuple[CAREamicsKiln, Configuration]
+        CAREamics model and configuration.
+
+    Raises
+    ------
+    ValueError
+        If the path is not a zip file.
+    """
+    path = Path(path)
+
+    if path.suffix != ".zip":
+        raise ValueError(f"Path must be a bioimage.io zip file, got {path}.")
+
+    # load description, this creates an unzipped folder next to the archive
+    model_desc = load_description(path)
+
+    # extract relative paths
+    weights_path, config_path = extract_model_path(model_desc)
+
+    # create folder path and absolute paths
+    unzip_path = get_unzip_path(path)
+    weights_path = unzip_path / weights_path
+    config_path = unzip_path / config_path
+
+    # load configuration
+    config = load_configuration(config_path)
+
+    # create careamics lightning module
+    model = CAREamicsKiln(algorithm_config=config.algorithm_config)
+
+    # load model state dictionary
+    _load_state_dict(model, weights_path)
+
+    return model, config

careamics/model_io/model_io_utils.py

@@ -0,0 +1,80 @@
+"""Utility functions to load pretrained models."""
+
+from pathlib import Path
+from typing import Tuple, Union
+
+from torch import load
+
+from careamics.config import Configuration
+from careamics.lightning_module import CAREamicsKiln
+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[CAREamicsKiln, 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[CAREamicsKiln, 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
+    checkpoint: dict = load(path)
+
+    # 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
+
+    model = CAREamicsKiln.load_from_checkpoint(path)
+
+    return model, Configuration(**cfg_dict)

careamics/models/__init__.py

@@ -1,4 +1,7 @@
 """Models package."""
 
-from .model_factory import create_model as create_model
+__all__ = ["model_factory", "UNet"]
+
+
+from .model_factory import model_factory
 from .unet import UNet as UNet