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

careamics/model_io/bioimage/model_description.py

@@ -4,6 +4,7 @@ from pathlib import Path
 from typing import List, Optional, Tuple, Union
 
 import numpy as np
+from bioimageio.spec._internal.io import resolve_and_extract
 from bioimageio.spec.model.v0_5 import (
     ArchitectureFromLibraryDescr,
     Author,
@@ -186,6 +187,7 @@ 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],
@@ -194,8 +196,9 @@ def create_model_description(
     careamics_version: str,
     config_path: Union[Path, str],
     env_path: Union[Path, str],
+    covers: list[Union[Path, str]],
     channel_names: Optional[List[str]] = None,
-    data_description: Optional[str] = None,
+    model_version: str = "0.1.0",
 ) -> ModelDescr:
     """Create model description.
 
@@ -207,6 +210,8 @@ def create_model_description(
         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]
@@ -223,10 +228,12 @@ def create_model_description(
         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.
-    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

@@ -6,8 +6,9 @@ from typing import List, Optional, Tuple, 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,
+    data_description: str,
     authors: List[dict],
     input_array: np.ndarray,
     output_array: np.ndarray,
+    covers: Optional[list[Union[Path, str]]] = None,
     channel_names: Optional[List[str]] = None,
-    data_description: Optional[str] = None,
+    model_version: str = "0.1.0",
 ) -> None:
     """Export the model to BioImage Model Zoo format.
 
@@ -110,16 +113,20 @@ def export_to_bmz(
         Model name.
     general_description : str
         General description of the model.
+    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.
+    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.
+    model_version : str, default="0.1.0"
+        Model version.
 
     Raises
     ------
@@ -161,16 +168,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 +191,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 +211,34 @@ def export_to_bmz(
 
 
 def load_from_bmz(
-    path: Union[Path, str]
+    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/models/layers.py

@@ -459,7 +459,8 @@ class MaxBlurPool(nn.Module):
         self.stride = stride
         self.max_pool_size = max_pool_size
         self.ceil_mode = ceil_mode
-        self.kernel = _get_pascal_kernel_nd(kernel_size, norm=True, dim=self.dim)
+        kernel = _get_pascal_kernel_nd(kernel_size, norm=True, dim=self.dim)
+        self.register_buffer("kernel", kernel, persistent=False)
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         """Forward pass of the function.
@@ -474,11 +475,12 @@ class MaxBlurPool(nn.Module):
         torch.Tensor
             Output tensor.
         """
-        self.kernel = torch.as_tensor(self.kernel, device=x.device, dtype=x.dtype)
+        kernel = self.kernel.to(dtype=x.dtype)
+        num_channels = int(x.size(1))
         if self.dim == 2:
             return _max_blur_pool_by_kernel2d(
                 x,
-                self.kernel.repeat((x.size(1), 1, 1, 1)),
+                kernel.repeat((num_channels, 1, 1, 1)),
                 self.stride,
                 self.max_pool_size,
                 self.ceil_mode,
@@ -486,7 +488,7 @@ class MaxBlurPool(nn.Module):
         else:
             return _max_blur_pool_by_kernel3d(
                 x,
-                self.kernel.repeat((x.size(1), 1, 1, 1, 1)),
+                kernel.repeat((num_channels, 1, 1, 1, 1)),
                 self.stride,
                 self.max_pool_size,
                 self.ceil_mode,