careamics 0.1.0rc5__py3-none-any.whl → 0.1.0rc7__py3-none-any.whl

careamics/model_io/bioimage/model_description.py

@@ -13,8 +13,8 @@ from bioimageio.spec.model.v0_5 import (
     ChannelAxis,
     EnvironmentFileDescr,
     FileDescr,
+    FixedZeroMeanUnitVarianceAlongAxisKwargs,
     FixedZeroMeanUnitVarianceDescr,
-    FixedZeroMeanUnitVarianceKwargs,
     Identifier,
     InputTensorDescr,
     ModelDescr,
@@ -134,44 +134,52 @@ def _create_inputs_ouputs(
     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
+    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_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
+    inv_means = []
+    inv_stds = []
+    if means and stds:
+        for mean, std in zip(means, stds):
+            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
+        return input_descr, output_descr
+    else:
+        raise ValueError("Mean and std cannot be None.")
 
 
 def create_model_description(
@@ -280,7 +288,7 @@ def create_model_description(
             "bioimageio": {
                 "test_kwargs": {
                     "pytorch_state_dict": {
-                        "decimals": 2,  # ...so we relax the constraints on the decimals
+                        "decimals": 0,  # ...so we relax the constraints on the decimals
                     }
                 }
             }

careamics/model_io/bmz_io.py

@@ -104,9 +104,9 @@ def export_to_bmz(
     authors : List[dict]
         Authors of the model.
     input_array : np.ndarray
-        Input array.
+        Input array, should not have been normalized.
     output_array : np.ndarray
-        Output array.
+        Output array, should have been denormalized.
     channel_names : Optional[List[str]], optional
         Channel names, by default None.
     data_description : Optional[str], optional
@@ -178,7 +178,7 @@ def export_to_bmz(
         )
 
         # test model description
-        summary: ValidationSummary = test_model(model_description, decimal=0)
+        summary: ValidationSummary = test_model(model_description, decimal=1)
         if summary.status == "failed":
             raise ValueError(f"Model description test failed: {summary}")
 

careamics/model_io/model_io_utils.py

@@ -3,7 +3,7 @@
 from pathlib import Path
 from typing import Tuple, Union
 
-from torch import load
+import torch
 
 from careamics.config import Configuration
 from careamics.lightning_module import CAREamicsModule
@@ -64,7 +64,10 @@ def _load_checkpoint(path: Union[Path, str]) -> Tuple[CAREamicsModule, Configura
         If the checkpoint file does not contain hyper parameters (configuration).
     """
     # load checkpoint
-    checkpoint: dict = load(path)
+    # 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:

careamics/models/activation.py

@@ -1,3 +1,5 @@
+"""Activations for CAREamics models."""
+
 from typing import Callable, Union
 
 import torch.nn as nn

careamics/models/layers.py

@@ -162,6 +162,18 @@ def _unpack_kernel_size(
     """Unpack kernel_size to a tuple of ints.
 
     Inspired by Kornia implementation. TODO: link
+
+    Parameters
+    ----------
+    kernel_size : Union[Tuple[int, ...], int]
+        Kernel size.
+    dim : int
+        Number of dimensions.
+
+    Returns
+    -------
+    Tuple[int, ...]
+        Kernel size tuple.
     """
     if isinstance(kernel_size, int):
         kernel_dims = tuple([kernel_size for _ in range(dim)])
@@ -173,7 +185,20 @@ def _unpack_kernel_size(
 def _compute_zero_padding(
     kernel_size: Union[Tuple[int, ...], int], dim: int
 ) -> Tuple[int, ...]:
-    """Utility function that computes zero padding tuple."""
+    """Utility function that computes zero padding tuple.
+
+    Parameters
+    ----------
+    kernel_size : Union[Tuple[int, ...], int]
+        Kernel size.
+    dim : int
+        Number of dimensions.
+
+    Returns
+    -------
+    Tuple[int, ...]
+        Zero padding tuple.
+    """
     kernel_dims = _unpack_kernel_size(kernel_size, dim)
     return tuple([(kd - 1) // 2 for kd in kernel_dims])
 
@@ -191,14 +216,19 @@ def get_pascal_kernel_1d(
 
     Parameters
     ----------
-    kernel_size: height and width of the kernel.
-    norm: if to normalize the kernel or not. Default: False.
-    device: tensor device
-    dtype: tensor dtype
+    kernel_size : int
+        Kernel size.
+    norm : bool
+        Normalize the kernel, by default False.
+    device : Optional[torch.device]
+        Device of the tensor, by default None.
+    dtype : Optional[torch.dtype]
+        Data type of the tensor, by default None.
 
     Returns
     -------
-    kernel shaped as :math:`(kernel_size,)`
+    torch.Tensor
+        Pascal kernel.
 
     Examples
     --------
@@ -245,19 +275,28 @@ def _get_pascal_kernel_nd(
 ) -> torch.Tensor:
     """Generate pascal filter kernel by kernel size.
 
+    If kernel_size is an integer the kernel will be shaped as (kernel_size, kernel_size)
+    otherwise the kernel will be shaped as kernel_size
+
     Inspired by Kornia implementation.
 
     Parameters
     ----------
-    kernel_size: height and width of the kernel.
-    norm: if to normalize the kernel or not. Default: True.
-    device: tensor device
-    dtype: tensor dtype
+    kernel_size : Union[Tuple[int, int], int]
+        Kernel size for the pascal kernel.
+    norm : bool
+        Normalize the kernel, by default True.
+    dim : int
+        Number of dimensions, by default 2.
+    device : Optional[torch.device]
+        Device of the tensor, by default None.
+    dtype : Optional[torch.dtype]
+        Data type of the tensor, by default None.
 
     Returns
     -------
-    if kernel_size is an integer the kernel will be shaped as (kernel_size, kernel_size)
-    otherwise the kernel will be shaped as kernel_size
+    torch.Tensor
+        Pascal kernel.
 
     Examples
     --------
@@ -303,6 +342,24 @@ def _max_blur_pool_by_kernel2d(
     """Compute max_blur_pool by a given :math:`CxC_(out, None)xNxN` kernel.
 
     Inspired by Kornia implementation.
+
+    Parameters
+    ----------
+    x : torch.Tensor
+        Input tensor.
+    kernel : torch.Tensor
+        Kernel tensor.
+    stride : int
+        Stride.
+    max_pool_size : int
+        Maximum pool size.
+    ceil_mode : bool
+        Ceil mode, by default False. Set to True to match output size of conv2d.
+
+    Returns
+    -------
+    torch.Tensor
+        Output tensor.
     """
     # compute local maxima
     x = F.max_pool2d(
@@ -323,6 +380,24 @@ def _max_blur_pool_by_kernel3d(
     """Compute max_blur_pool by a given :math:`CxC_(out, None)xNxNxN` kernel.
 
     Inspired by Kornia implementation.
+
+    Parameters
+    ----------
+    x : torch.Tensor
+        Input tensor.
+    kernel : torch.Tensor
+        Kernel tensor.
+    stride : int
+        Stride.
+    max_pool_size : int
+        Maximum pool size.
+    ceil_mode : bool
+        Ceil mode, by default False. Set to True to match output size of conv2d.
+
+    Returns
+    -------
+    torch.Tensor
+        Output tensor.
     """
     # compute local maxima
     x = F.max_pool3d(
@@ -343,21 +418,16 @@ class MaxBlurPool(nn.Module):
 
     Parameters
     ----------
-    dim: int
-        Toggles between 2D and 3D
-    kernel_size: Union[Tuple[int, int], int]
+    dim : int
+        Toggles between 2D and 3D.
+    kernel_size : Union[Tuple[int, int], int]
         Kernel size for max pooling.
-    stride: int
+    stride : int
         Stride for pooling.
-    max_pool_size: int
+    max_pool_size : int
         Max kernel size for max pooling.
-    ceil_mode: bool
-        Should be true to match output size of conv2d with same kernel size.
-
-    Returns
-    -------
-    torch.Tensor
-        The pooled and blurred tensor.
+    ceil_mode : bool
+        Ceil mode, by default False. Set to True to match output size of conv2d.
     """
 
     def __init__(
@@ -368,6 +438,21 @@ class MaxBlurPool(nn.Module):
         max_pool_size: int = 2,
         ceil_mode: bool = False,
     ) -> None:
+        """Constructor.
+
+        Parameters
+        ----------
+        dim : int
+            Dimension of the convolution.
+        kernel_size : Union[Tuple[int, int], int]
+            Kernel size for max pooling.
+        stride : int, optional
+            Stride, by default 2.
+        max_pool_size : int, optional
+            Maximum pool size, by default 2.
+        ceil_mode : bool, optional
+            Ceil mode, by default False. Set to True to match output size of conv2d.
+        """
         super().__init__()
         self.dim = dim
         self.kernel_size = kernel_size
@@ -377,7 +462,18 @@ class MaxBlurPool(nn.Module):
         self.kernel = _get_pascal_kernel_nd(kernel_size, norm=True, dim=self.dim)
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
-        """Forward pass of the function."""
+        """Forward pass of the function.
+
+        Parameters
+        ----------
+        x : torch.Tensor
+            Input tensor.
+
+        Returns
+        -------
+        torch.Tensor
+            Output tensor.
+        """
         self.kernel = torch.as_tensor(self.kernel, device=x.device, dtype=x.dtype)
         if self.dim == 2:
             return _max_blur_pool_by_kernel2d(