careamics 0.1.0rc6__py3-none-any.whl → 0.1.0rc8__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

@@ -12,7 +12,7 @@ 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 CAREamicsModule
+from careamics.lightning.lightning_module import CAREamicsModule
 
 from .bioimage import (
     create_env_text,
@@ -178,7 +178,7 @@ def export_to_bmz(
         )
 
         # test model description
-        summary: ValidationSummary = test_model(model_description, decimal=2)
+        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,10 +3,10 @@
 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
+from careamics.lightning.lightning_module import CAREamicsModule
 from careamics.model_io.bmz_io import load_from_bmz
 from careamics.utils import check_path_exists
 
@@ -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: