careamics 0.0.19__py3-none-any.whl

careamics/model_io/bioimage/_readme_factory.py

@@ -0,0 +1,113 @@
+"""Functions used to create a README.md file for BMZ export."""
+
+from pathlib import Path
+
+import yaml
+
+from careamics.config import Configuration
+from careamics.utils import cwd, get_careamics_home
+
+
+def _yaml_block(yaml_str: str) -> str:
+    """Return a markdown code block with a yaml string.
+
+    Parameters
+    ----------
+    yaml_str : str
+        YAML string.
+
+    Returns
+    -------
+    str
+        Markdown code block with the YAML string.
+    """
+    return f"```yaml\n{yaml_str}\n```"
+
+
+def readme_factory(
+    config: Configuration,
+    careamics_version: str,
+    data_description: str,
+) -> Path:
+    """Create a README file for the model.
+
+    `data_description` can be used to add more information about the content of the
+    data the model was trained on.
+
+    Parameters
+    ----------
+    config : Configuration
+        CAREamics configuration.
+    careamics_version : str
+        CAREamics version.
+    data_description : str
+        Description of the data.
+
+    Returns
+    -------
+    Path
+        Path to the README file.
+    """
+    # create file
+    # TODO use tempfile as in the bmz_io module
+    with cwd(get_careamics_home()):
+        readme = Path("README.md")
+        readme.touch()
+
+        # algorithm pretty name
+        algorithm_flavour = config.get_algorithm_friendly_name()
+        algorithm_pretty_name = algorithm_flavour + " - CAREamics"
+
+        description = [f"# {algorithm_pretty_name}\n\n"]
+
+        # data description
+        description.append("## Data description\n\n")
+        description.append(data_description)
+        description.append("\n\n")
+
+        # algorithm description
+        description.append("## Algorithm description:\n\n")
+        description.append(config.get_algorithm_description())
+        description.append("\n\n")
+
+        # configuration description
+        description.append("## Configuration\n\n")
+
+        description.append(
+            f"{algorithm_flavour} was trained using CAREamics (version "
+            f"{careamics_version}) using the following configuration:\n\n"
+        )
+
+        description.append(_yaml_block(yaml.dump(config.model_dump(exclude_none=True))))
+        description.append("\n\n")
+
+        # validation
+        description.append("# Validation\n\n")
+
+        description.append(
+            "In order to validate the model, we encourage users to acquire a "
+            "test dataset with ground-truth data. Comparing the ground-truth data "
+            "with the prediction allows unbiased evaluation of the model performances. "
+            "This can be done for instance by using metrics such as PSNR, SSIM, or"
+            "MicroSSIM. In the absence of ground-truth, inspecting the residual image "
+            "(difference between input and predicted image) can be helpful to identify "
+            "whether real signal is removed from the input image.\n\n"
+        )
+
+        # references
+        reference = config.get_algorithm_references()
+        if reference != "":
+            description.append("## References\n\n")
+            description.append(reference)
+            description.append("\n\n")
+
+        # links
+        description.append(
+            "# Links\n\n"
+            "- [CAREamics repository](https://github.com/CAREamics/careamics)\n"
+            "- [CAREamics documentation](https://careamics.github.io/)\n"
+        )
+
+        readme.write_text("".join(description))
+
+        return readme.absolute()

careamics/model_io/bioimage/bioimage_utils.py

@@ -0,0 +1,56 @@
+"""Bioimage.io utils."""
+
+from pathlib import Path
+from typing import Union
+
+from careamics.utils.version import get_careamics_version
+
+
+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, torchvision_version: str) -> str:
+    """Create environment yaml content for the bioimage model.
+
+    This installs an environment with the specified pytorch version and the latest
+    changes to careamics.
+
+    Parameters
+    ----------
+    pytorch_version : str
+        Pytorch version.
+    torchvision_version : str
+        Torchvision version.
+
+    Returns
+    -------
+    str
+        Environment text.
+    """
+    env = (
+        f"name: careamics\n"
+        f"dependencies:\n"
+        f"  - python=3.12\n"
+        f"  - pip\n"
+        f"  - pip:\n"
+        f"    - torch=={pytorch_version}\n"
+        f"    - torchvision=={torchvision_version}\n"
+        f"    - careamics=={get_careamics_version()}"
+    )
+
+    return env

careamics/model_io/bioimage/cover_factory.py

@@ -0,0 +1,171 @@
+"""Convenience function to create covers for the BMZ."""
+
+from pathlib import Path
+
+import numpy as np
+from numpy.typing import NDArray
+from PIL import Image
+
+color_palette = np.array(
+    [
+        np.array([255, 195, 0]),  # grey
+        np.array([189, 226, 240]),
+        np.array([96, 60, 76]),
+        np.array([193, 225, 193]),
+    ]
+)
+
+
+def _get_norm_slice(array: NDArray) -> NDArray:
+    """Get the normalized middle slice of a 4D or 5D array (SC(Z)YX).
+
+    Parameters
+    ----------
+    array : NDArray
+        Array from which to get the middle slice.
+
+    Returns
+    -------
+    NDArray
+        Normalized middle slice of the input array.
+    """
+    if array.ndim not in (4, 5):
+        raise ValueError("Array must be 4D or 5D.")
+
+    channels = array.shape[1] > 1
+    z_stack = array.ndim == 5
+
+    # get slice
+    if z_stack:
+        array_slice = array[0, :, array.shape[2] // 2, ...]
+    else:
+        array_slice = array[0, ...]
+
+    # channels
+    if channels:
+        array_slice = np.moveaxis(array_slice, 0, -1)
+    else:
+        array_slice = array_slice[0, ...]
+
+    # normalize
+    array_slice = (
+        255
+        * (array_slice - array_slice.min())
+        / (array_slice.max() - array_slice.min())
+    )
+
+    return array_slice.astype(np.uint8)
+
+
+def _four_channel_image(array: NDArray) -> Image:
+    """Convert 4-channel array to Image.
+
+    Parameters
+    ----------
+    array : NDArray
+        Normalized array to convert.
+
+    Returns
+    -------
+    Image
+        Converted array.
+    """
+    colors = color_palette[np.newaxis, np.newaxis, :, :]
+    four_c_array = np.sum(array[..., :4, np.newaxis] * colors, axis=-2).astype(np.uint8)
+
+    return Image.fromarray(four_c_array).convert("RGB")
+
+
+def _convert_to_image(original_shape: tuple[int, ...], array: NDArray) -> Image:
+    """Convert to Image.
+
+    Parameters
+    ----------
+    original_shape : tuple
+        Original shape of the array.
+    array : NDArray
+        Normalized array to convert.
+
+    Returns
+    -------
+    Image
+        Converted array.
+    """
+    n_channels = original_shape[1]
+
+    if n_channels > 1:
+        if n_channels == 3:
+            return Image.fromarray(array).convert("RGB")
+        elif n_channels == 2:
+            # add an empty channel to the numpy array
+            array = np.concatenate([np.zeros_like(array[..., 0:1]), array], axis=-1)
+
+            return Image.fromarray(array).convert("RGB")
+        else:  # more than 4
+            return _four_channel_image(array[..., :4])
+    else:
+        return Image.fromarray(array).convert("L").convert("RGB")
+
+
+def create_cover(directory: Path, array_in: NDArray, array_out: NDArray) -> Path:
+    """Create a cover image from input and output arrays.
+
+    Input and output arrays are expected to be SC(Z)YX. For images with a Z
+    dimension, the middle slice is taken.
+
+    Parameters
+    ----------
+    directory : Path
+        Directory in which to save the cover.
+    array_in : numpy.ndarray
+        Array from which to create the cover image.
+    array_out : numpy.ndarray
+        Array from which to create the cover image.
+
+    Returns
+    -------
+    Path
+        Path to the saved cover image.
+    """
+    # extract slice and normalize arrays
+    slice_in = _get_norm_slice(array_in)
+    slice_out = _get_norm_slice(array_out)
+
+    horizontal_split = slice_in.shape[-1] == slice_out.shape[-1]
+    if not horizontal_split:
+        if slice_in.shape[-2] != slice_out.shape[-2]:
+            raise ValueError("Input and output arrays have different shapes.")
+
+    # convert to Image
+    image_in = _convert_to_image(array_in.shape, slice_in)
+    image_out = _convert_to_image(array_out.shape, slice_out)
+
+    # split horizontally or vertically
+    if horizontal_split:
+        width = image_in.width // 2
+
+        cover = Image.new("RGB", (image_in.width, image_in.height))
+        cover.paste(image_in.crop((0, 0, width, image_in.height)), (0, 0))
+        cover.paste(
+            image_out.crop(
+                (image_in.width - width, 0, image_in.width, image_in.height)
+            ),
+            (width, 0),
+        )
+    else:
+        height = image_in.height // 2
+
+        cover = Image.new("RGB", (image_in.width, image_in.height))
+        cover.paste(image_in.crop((0, 0, image_in.width, height)), (0, 0))
+        cover.paste(
+            image_out.crop(
+                (0, image_in.height - height, image_in.width, image_in.height)
+            ),
+            (0, height),
+        )
+
+    # save
+    cover_path = directory / "cover.png"
+    cover.save(cover_path)
+
+    return cover_path

careamics/model_io/bioimage/model_description.py

@@ -0,0 +1,341 @@
+"""Module use to build BMZ model description."""
+
+from pathlib import Path
+from typing import Union
+
+import numpy as np
+from bioimageio.spec._internal.io import extract
+from bioimageio.spec.model.v0_5 import (
+    ArchitectureFromLibraryDescr,
+    Author,
+    AxisBase,
+    AxisId,
+    BatchAxis,
+    ChannelAxis,
+    FileDescr,
+    FixedZeroMeanUnitVarianceAlongAxisKwargs,
+    FixedZeroMeanUnitVarianceDescr,
+    Identifier,
+    InputTensorDescr,
+    ModelDescr,
+    OutputTensorDescr,
+    PytorchStateDictWeightsDescr,
+    SpaceInputAxis,
+    SpaceOutputAxis,
+    TensorId,
+    Version,
+    WeightsDescr,
+)
+
+from careamics.config import Configuration, DataConfig
+
+from ._readme_factory import readme_factory
+
+
+def _create_axes(
+    array: np.ndarray,
+    data_config: DataConfig,
+    channel_names: list[str] | None = 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: DataConfig,
+    input_path: Union[Path, str],
+    output_path: Union[Path, str],
+    channel_names: list[str] | None = 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.image_means is not None, "Mean cannot be None."
+    assert data_config.image_means is not None, "Std cannot be None."
+    means = data_config.image_means
+    stds = data_config.image_stds
+
+    # 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_means = []
+    inv_stds = []
+    if means and stds:
+        for mean, std in zip(means, stds, strict=False):
+            inv_means.append(-mean / (std + eps))
+            inv_stds.append(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=FixedZeroMeanUnitVarianceAlongAxisKwargs(
+                        mean=means, std=stds, axis="channel"
+                    )
+                )
+            ],
+        )
+        output_descr = OutputTensorDescr(
+            id=TensorId("prediction"),
+            axes=output_axes,
+            test_tensor=FileDescr(source=output_path),
+            postprocessing=[
+                FixedZeroMeanUnitVarianceDescr(
+                    kwargs=FixedZeroMeanUnitVarianceAlongAxisKwargs(  # invert norm
+                        mean=inv_means, std=inv_stds, axis="channel"
+                    )
+                )
+            ],
+        )
+
+        return input_descr, output_descr
+    else:
+        raise ValueError("Mean and std cannot be None.")
+
+
+def create_model_description(
+    config: Configuration,
+    name: str,
+    general_description: str,
+    data_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],
+    covers: list[Union[Path, str]],
+    channel_names: list[str] | None = None,
+    model_version: str = "0.1.0",
+) -> ModelDescr:
+    """Create model description.
+
+    Parameters
+    ----------
+    config : Configuration
+        CAREamics configuration.
+    name : str
+        Name of the model.
+    general_description : str
+        General description of the model.
+    data_description : str
+        Description of the data the model was trained on.
+    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.
+    covers : list of pathlib.Path or str
+        Paths to cover images.
+    channel_names : Optional[list[str]], optional
+        Channel names, by default None.
+    model_version : str, default "0.1.0"
+        Model version.
+
+    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.unet",
+        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=FileDescr(source=Path(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",
+        config={
+            "bioimageio": {
+                "test_kwargs": {
+                    "pytorch_state_dict": {
+                        "absolute_tolerance": 1e-2,
+                        "relative_tolerance": 1e-2,
+                    }
+                }
+            }
+        },
+        version=model_version,
+        weights=weights_descr,
+        attachments=[FileDescr(source=config_path)],
+        cite=config.get_algorithm_citations(),
+        covers=covers,
+    )
+
+    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 of (path, path)
+        Weights and configuration paths.
+    """
+    if model_desc.weights.pytorch_state_dict is None:
+        raise ValueError("No model weights found in model description.")
+
+    # extract the zip model and return the directory
+    model_dir = extract(model_desc.root)
+
+    weights_path = model_dir.joinpath(model_desc.weights.pytorch_state_dict.source.path)
+
+    for file in model_desc.attachments:
+        file_path = file.source if isinstance(file.source, Path) else file.source.path
+        if file_path is None:
+            continue
+        file_path = Path(file_path)
+        if file_path.name == "careamics.yaml":
+            config_path = model_dir.joinpath(file.source.path)
+            break
+    else:
+        raise ValueError("Configuration file not found.")
+
+    return weights_path, config_path