careamics 0.0.4.2__py3-none-any.whl → 0.0.6__py3-none-any.whl

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

@@ -1,9 +1,10 @@
 """Module use to build BMZ model description."""
 
 from pathlib import Path
-from typing import List, Optional, Tuple, Union
+from typing import Optional, Union
 
 import numpy as np
+from bioimageio.spec._internal.io import resolve_and_extract
 from bioimageio.spec.model.v0_5 import (
     ArchitectureFromLibraryDescr,
     Author,
@@ -27,17 +28,17 @@ from bioimageio.spec.model.v0_5 import (
     WeightsDescr,
 )
 
-from careamics.config import Configuration, DataConfig
+from careamics.config import Configuration, GeneralDataConfig
 
 from ._readme_factory import readme_factory
 
 
 def _create_axes(
     array: np.ndarray,
-    data_config: DataConfig,
-    channel_names: Optional[List[str]] = None,
+    data_config: GeneralDataConfig,
+    channel_names: Optional[list[str]] = None,
     is_input: bool = True,
-) -> List[AxisBase]:
+) -> list[AxisBase]:
     """Create axes description.
 
     Array shape is expected to be SC(Z)YX.
@@ -48,15 +49,15 @@ def _create_axes(
         Array.
     data_config : DataModel
         CAREamics data configuration.
-    channel_names : Optional[List[str]], optional
+    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.
+    list[AxisBase]
+        list of axes description.
 
     Raises
     ------
@@ -101,11 +102,11 @@ def _create_axes(
 def _create_inputs_ouputs(
     input_array: np.ndarray,
     output_array: np.ndarray,
-    data_config: DataConfig,
+    data_config: GeneralDataConfig,
     input_path: Union[Path, str],
     output_path: Union[Path, str],
-    channel_names: Optional[List[str]] = None,
-) -> Tuple[InputTensorDescr, OutputTensorDescr]:
+    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.
@@ -122,12 +123,12 @@ def _create_inputs_ouputs(
         Path to input .npy file.
     output_path : Union[Path, str]
         Path to output .npy file.
-    channel_names : Optional[List[str]], optional
+    channel_names : Optional[list[str]], optional
         Channel names, by default None.
 
     Returns
     -------
-    Tuple[InputTensorDescr, OutputTensorDescr]
+    tuple[InputTensorDescr, OutputTensorDescr]
         Input and output tensor descriptions.
     """
     input_axes = _create_axes(input_array, data_config, channel_names)
@@ -186,7 +187,8 @@ def create_model_description(
     config: Configuration,
     name: str,
     general_description: str,
-    authors: List[Author],
+    data_description: str,
+    authors: list[Author],
     inputs: Union[Path, str],
     outputs: Union[Path, str],
     weights_path: Union[Path, str],
@@ -194,8 +196,9 @@ def create_model_description(
     careamics_version: str,
     config_path: Union[Path, str],
     env_path: Union[Path, str],
-    channel_names: Optional[List[str]] = None,
-    data_description: Optional[str] = None,
+    covers: list[Union[Path, str]],
+    channel_names: Optional[list[str]] = None,
+    model_version: str = "0.1.0",
 ) -> ModelDescr:
     """Create model description.
 
@@ -207,7 +210,9 @@ def create_model_description(
         Name of the model.
     general_description : str
         General description of the model.
-    authors : List[Author]
+    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.
@@ -223,10 +228,12 @@ def create_model_description(
         Path to model configuration.
     env_path : Union[Path, str]
         Path to environment file.
-    channel_names : Optional[List[str]], optional
+    covers : list of pathlib.Path or str
+        Paths to cover images.
+    channel_names : Optional[list[str]], optional
         Channel names, by default None.
-    data_description : Optional[str], optional
-        Description of the data, by default None.
+    model_version : str, default "0.1.0"
+        Model version.
 
     Returns
     -------
@@ -280,16 +287,27 @@ def create_model_description(
             "https://careamics.github.io/latest/",
         ],
         license="BSD-3-Clause",
-        version="0.1.0",
+        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]:
+def extract_model_path(model_desc: ModelDescr) -> tuple[Path, Path]:
     """Return the relative path to the weights and configuration files.
 
     Parameters
@@ -299,20 +317,24 @@ def extract_model_path(model_desc: ModelDescr) -> Tuple[Path, Path]:
 
     Returns
     -------
-    Tuple[Path, Path]
+    tuple of (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
+    if model_desc.weights.pytorch_state_dict is None:
+        raise ValueError("No model weights found in model description.")
+    weights_path = resolve_and_extract(
+        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 = resolve_and_extract(file.source).path
+            break
     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.")
+        raise ValueError("Configuration file not found.")
 
     return weights_path, config_path

careamics/model_io/bmz_io.py

@@ -2,12 +2,13 @@
 
 import tempfile
 from pathlib import Path
-from typing import List, Optional, Tuple, Union
+from typing import Optional, Union
 
 import numpy as np
 import pkg_resources
-from bioimageio.core import load_description, test_model
+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
@@ -20,8 +21,8 @@ from .bioimage import (
     create_env_text,
     create_model_description,
     extract_model_path,
-    get_unzip_path,
 )
+from .bioimage.cover_factory import create_cover
 
 
 def _export_state_dict(
@@ -85,11 +86,13 @@ def export_to_bmz(
     path_to_archive: Union[Path, str],
     model_name: str,
     general_description: str,
-    authors: List[dict],
+    data_description: str,
+    authors: list[dict],
     input_array: np.ndarray,
     output_array: np.ndarray,
-    channel_names: Optional[List[str]] = None,
-    data_description: Optional[str] = None,
+    covers: Optional[list[Union[Path, str]]] = None,
+    channel_names: Optional[list[str]] = None,
+    model_version: str = "0.1.0",
 ) -> None:
     """Export the model to BioImage Model Zoo format.
 
@@ -110,30 +113,23 @@ def export_to_bmz(
         Model name.
     general_description : str
         General description of the model.
-    authors : List[dict]
+    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.
-    channel_names : Optional[List[str]], optional
+    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.
-    data_description : Optional[str], optional
-        Description of the data, by default None.
-
-    Raises
-    ------
-    ValueError
-        If the model is a Custom model.
+    model_version : str, default="0.1.0"
+        Model version.
     """
     path_to_archive = Path(path_to_archive)
 
-    # 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."
-        )
-
     if path_to_archive.suffix != ".zip":
         raise ValueError(
             f"Path to archive must point to a zip file, got {path_to_archive}."
@@ -161,16 +157,21 @@ def export_to_bmz(
         np.save(outputs, output_array)
 
         # export configuration
-        config_path = save_configuration(config, temp_path)
+        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,
@@ -179,12 +180,18 @@ def export_to_bmz(
             careamics_version=careamics_version,
             config_path=config_path,
             env_path=env_path,
+            covers=covers,
             channel_names=channel_names,
-            data_description=data_description,
+            model_version=model_version,
         )
 
         # test model description
-        summary: ValidationSummary = test_model(model_description)
+        test_kwargs = (
+            model_description.config.get("bioimageio", {})
+            .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}")
 
@@ -193,41 +200,34 @@ def export_to_bmz(
 
 
 def load_from_bmz(
-    path: Union[Path, str]
-) -> Tuple[Union[FCNModule, VAEModule], Configuration]:
+    path: Union[Path, str, HttpUrl]
+) -> tuple[Union[FCNModule, VAEModule], Configuration]:
     """Load a model from a BioImage Model Zoo archive.
 
     Parameters
     ----------
-    path : Union[Path, str]
-        Path to the BioImage Model Zoo archive.
+    path : Path, str or HttpUrl
+        Path to the BioImage Model Zoo archive. A Http URL must point to a downloadable
+        location.
 
     Returns
     -------
-    Tuple[CAREamicsKiln, Configuration]
-        CAREamics model and configuration.
+    FCNModel or VAEModel
+        The loaded CAREamics model.
+    Configuration
+        The loaded CAREamics 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)
+    model_desc = load_model_description(path)
 
-    # extract relative paths
+    # extract 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)
 

careamics/model_io/model_io_utils.py

@@ -1,11 +1,11 @@
 """Utility functions to load pretrained models."""
 
 from pathlib import Path
-from typing import Tuple, Union
+from typing import Union
 
 import torch
 
-from careamics.config import Configuration
+from careamics.config import Configuration, configuration_factory
 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
@@ -13,7 +13,7 @@ from careamics.utils import check_path_exists
 
 def load_pretrained(
     path: Union[Path, str]
-) -> Tuple[Union[FCNModule, VAEModule], Configuration]:
+) -> tuple[Union[FCNModule, VAEModule], Configuration]:
     """
     Load a pretrained model from a checkpoint or a BioImage Model Zoo model.
 
@@ -26,8 +26,8 @@ def load_pretrained(
 
     Returns
     -------
-    Tuple[CAREamicsKiln, Configuration]
-        Tuple of CAREamics model and its configuration.
+    tuple[CAREamicsKiln, Configuration]
+        tuple of CAREamics model and its configuration.
 
     Raises
     ------
@@ -48,7 +48,7 @@ def load_pretrained(
 
 def _load_checkpoint(
     path: Union[Path, str]
-) -> Tuple[Union[FCNModule, VAEModule], Configuration]:
+) -> tuple[Union[FCNModule, VAEModule], Configuration]:
     """
     Load a model from a checkpoint and return both model and configuration.
 
@@ -59,8 +59,8 @@ def _load_checkpoint(
 
     Returns
     -------
-    Tuple[CAREamicsKiln, Configuration]
-        Tuple of CAREamics model and its configuration.
+    tuple[CAREamicsKiln, Configuration]
+        tuple of CAREamics model and its configuration.
 
     Raises
     ------
@@ -92,4 +92,4 @@ def _load_checkpoint(
             f"{cfg_dict['algorithm_config']['model']['architecture']}"
         )
 
-    return model, Configuration(**cfg_dict)
+    return model, configuration_factory(cfg_dict)