careamics 0.1.0rc1__py3-none-any.whl → 0.1.0rc2__py3-none-any.whl

careamics/bioimage/__init__.py

@@ -1,7 +1,7 @@
 """Provide utilities for exporting models to BioImage model zoo."""
 
 __all__ = [
-    "build_zip_model",
+    "save_bioimage_model",
     "import_bioimage_model",
     "get_default_model_specs",
     "PYTORCH_STATE_DICT",
@@ -9,7 +9,7 @@ __all__ = [
 
 from .io import (
     PYTORCH_STATE_DICT,
-    build_zip_model,
-    get_default_model_specs,
     import_bioimage_model,
+    save_bioimage_model,
 )
+from .rdf import get_default_model_specs

careamics/bioimage/io.py

@@ -5,125 +5,20 @@ from typing import Union
 import torch
 from bioimageio.core import load_resource_description
 from bioimageio.core.build_spec import build_model
-from bioimageio.spec.model.raw_nodes import Model
 
-from careamics.config.config import (
-    Configuration,
-)
+from careamics.config.config import Configuration
+from careamics.utils.context import cwd
 
 PYTORCH_STATE_DICT = "pytorch_state_dict"
 
 
-def _get_model_doc(name: str) -> str:
-    """
-    Return markdown documentation path for the provided model.
-
-    Parameters
-    ----------
-    name : str
-        Model's name.
-
-    Returns
-    -------
-    str
-        Path to the model's markdown documentation.
-
-    Raises
-    ------
-    FileNotFoundError
-        If the documentation file was not found.
-    """
-    doc = Path(__file__).parent.joinpath("docs").joinpath(f"{name}.md")
-    if doc.exists():
-        return str(doc.absolute())
-    else:
-        raise FileNotFoundError(f"Documentation for {name} was not found.")
-
-
-def get_default_model_specs(
-    name: str, mean: float, std: float, is_3D: bool = False
-) -> dict:
-    """
-    Return the default bioimage.io specs for the provided model's name.
-
-    Currently only supports `n2v` model.
-
-    Parameters
-    ----------
-    name : str
-        Algorithm's name.
-    mean : float
-        Mean of the dataset.
-    std : float
-        Std of the dataset.
-    is_3D : bool, optional
-        Whether the model is 3D or not, by default False.
-
-    Returns
-    -------
-    dict
-        Model specs compatible with bioimage.io export.
-    """
-    rdf = {
-        "name": "Noise2Void",
-        "description": "Self-supervised denoising.",
-        "license": "BSD-3-Clause",
-        "authors": [
-            {"name": "Alexander Krull"},
-            {"name": "Tim-Oliver Buchholz"},
-            {"name": "Florian Jug"},
-        ],
-        "cite": [
-            {
-                "doi": "10.48550/arXiv.1811.10980",
-                "text": 'A. Krull, T.-O. Buchholz and F. Jug, "Noise2Void - Learning '
-                'Denoising From Single Noisy Images," 2019 IEEE/CVF '
-                "Conference on Computer Vision and Pattern Recognition "
-                "(CVPR), 2019, pp. 2124-2132",
-            }
-        ],
-        # "input_axes": ["bcyx"], <- overriden in save_as_bioimage
-        "preprocessing": [  # for multiple inputs
-            [  # multiple processes per input
-                {
-                    "kwargs": {
-                        "axes": "zyx" if is_3D else "yx",
-                        "mean": [mean],
-                        "mode": "fixed",
-                        "std": [std],
-                    },
-                    "name": "zero_mean_unit_variance",
-                }
-            ]
-        ],
-        # "output_axes": ["bcyx"], <- overriden in save_as_bioimage
-        "postprocessing": [  # for multiple outputs
-            [  # multiple processes per input
-                {
-                    "kwargs": {
-                        "axes": "zyx" if is_3D else "yx",
-                        "gain": [std],
-                        "offset": [mean],
-                    },
-                    "name": "scale_linear",
-                }
-            ]
-        ],
-        "tags": ["unet", "denoising", "Noise2Void", "tensorflow", "napari"],
-    }
-
-    rdf["documentation"] = _get_model_doc(name)
-
-    return rdf
-
-
-def build_zip_model(
+def save_bioimage_model(
     path: Union[str, Path],
     config: Configuration,
-    model_specs: dict,
-) -> Model:
+    specs: dict,
+) -> None:
     """
-    Build bioimage model zip file from model specification data.
+    Build bioimage model zip file from model RDF data.
 
     Parameters
     ----------
@@ -131,71 +26,86 @@ def build_zip_model(
         Path to the model zip file.
     config : Configuration
         Configuration object.
-    model_specs : dict
-        Model specification data.
-
-    Returns
-    -------
-    Model
-        Bioimage model object.
+    specs : dict
+        Model RDF dict.
     """
     workdir = config.working_directory
 
-    # load best checkpoint
-    checkpoint_path = workdir.joinpath(f"{config.experiment_name}_best.pth").absolute()
-    checkpoint = torch.load(checkpoint_path, map_location="cpu")
-
-    # save chekpoint entries in separate files
-    weight_path = workdir.joinpath("model_weights.pth")
-    torch.save(checkpoint["model_state_dict"], weight_path)
-
-    optim_path = workdir.joinpath("optim.pth")
-    torch.save(checkpoint["optimizer_state_dict"], optim_path)
-
-    scheduler_path = workdir.joinpath("scheduler.pth")
-    torch.save(checkpoint["scheduler_state_dict"], scheduler_path)
-
-    grad_path = workdir.joinpath("grad.pth")
-    torch.save(checkpoint["grad_scaler_state_dict"], grad_path)
-
-    config_path = workdir.joinpath("config.pth")
-    torch.save(config.model_dump(), config_path)
-
-    # Create attachments
-    attachments = [
-        str(optim_path),
-        str(scheduler_path),
-        str(grad_path),
-        str(config_path),
-    ]
-
-    model_specs.update(
-        {
-            "weight_type": PYTORCH_STATE_DICT,
-            "weight_uri": str(weight_path),
-            "attachments": {"files": attachments},
-        }
-    )
+    # temporary folder
+    temp_folder = Path.home().joinpath(".careamics", "bmz_tmp")
+    temp_folder.mkdir(exist_ok=True, parents=True)
+
+    # change working directory to the temp folder
+    with cwd(temp_folder):
+        # load best checkpoint
+        checkpoint_path = workdir.joinpath(
+            f"{config.experiment_name}_best.pth"
+        ).absolute()
+        checkpoint = torch.load(checkpoint_path, map_location="cpu")
+
+        # save chekpoint entries in separate files
+        weight_path = Path("model_weights.pth")
+        torch.save(checkpoint["model_state_dict"], weight_path)
+
+        optim_path = Path("optim.pth")
+        torch.save(checkpoint["optimizer_state_dict"], optim_path)
+
+        scheduler_path = Path("scheduler.pth")
+        torch.save(checkpoint["scheduler_state_dict"], scheduler_path)
+
+        grad_path = Path("grad.pth")
+        torch.save(checkpoint["grad_scaler_state_dict"], grad_path)
+
+        config_path = Path("config.pth")
+        torch.save(config.model_dump(), config_path)
+
+        # create attachments
+        attachments = [
+            str(optim_path),
+            str(scheduler_path),
+            str(grad_path),
+            str(config_path),
+        ]
+
+        # create requirements file
+        requirements = Path("requirements.txt")
+        with open(requirements, "w") as f:
+            f.write("git+https://github.com/CAREamics/careamics.git")
+
+        algo_config = config.algorithm
+        specs.update(
+            {
+                "weight_type": PYTORCH_STATE_DICT,
+                "weight_uri": str(weight_path),
+                "architecture": "careamics.models.unet.UNet",
+                "pytorch_version": torch.__version__,
+                "model_kwargs": {
+                    "conv_dim": algo_config.get_conv_dim(),
+                    "depth": algo_config.model_parameters.depth,
+                    "num_channels_init": algo_config.model_parameters.num_channels_init,
+                },
+                "dependencies": "pip:" + str(requirements),
+                "attachments": {"files": attachments},
+            }
+        )
 
-    if config.algorithm.is_3D:
-        model_specs["tags"].append("3D")
-    else:
-        model_specs["tags"].append("2D")
+        if config.algorithm.is_3D:
+            specs["tags"].append("3D")
+        else:
+            specs["tags"].append("2D")
 
-    # build model zip
-    raw_model = build_model(
-        output_path=Path(path).absolute(),
-        **model_specs,
-    )
+        # build model zip
+        build_model(
+            output_path=Path(path).absolute(),
+            **specs,
+        )
 
-    # remove the temporary files
-    weight_path.unlink()
-    optim_path.unlink()
-    scheduler_path.unlink()
-    grad_path.unlink()
-    config_path.unlink()
+        # remove temporary files
+        for file in temp_folder.glob("*"):
+            file.unlink()
 
-    return raw_model
+    # delete temporary folder
+    temp_folder.rmdir()
 
 
 def import_bioimage_model(model_path: Union[str, Path]) -> Path:
@@ -219,11 +129,12 @@ def import_bioimage_model(model_path: Union[str, Path]) -> Path:
     FileNotFoundError
         If the checkpoint file was not found.
     """
-    if isinstance(model_path, str):
-        model_path = Path(model_path)
+    model_path = Path(model_path)
+
     # check the model extension (should be a zip file).
     if model_path.suffix != ".zip":
         raise ValueError("Invalid model format. Expected bioimage model zip file.")
+
     # load the model
     rdf = load_resource_description(model_path)
 

careamics/bioimage/rdf.py

@@ -0,0 +1,105 @@
+"""RDF related methods."""
+from pathlib import Path
+
+
+def _get_model_doc(name: str) -> str:
+    """
+    Return markdown documentation path for the provided model.
+
+    Parameters
+    ----------
+    name : str
+        Model's name.
+
+    Returns
+    -------
+    str
+        Path to the model's markdown documentation.
+
+    Raises
+    ------
+    FileNotFoundError
+        If the documentation file was not found.
+    """
+    doc = Path(__file__).parent.joinpath("docs").joinpath(f"{name}.md")
+    if doc.exists():
+        return str(doc.absolute())
+    else:
+        raise FileNotFoundError(f"Documentation for {name} was not found.")
+
+
+def get_default_model_specs(
+    name: str, mean: float, std: float, is_3D: bool = False
+) -> dict:
+    """
+    Return the default bioimage.io specs for the provided model's name.
+
+    Currently only supports `Noise2Void` model.
+
+    Parameters
+    ----------
+    name : str
+        Algorithm's name.
+    mean : float
+        Mean of the dataset.
+    std : float
+        Std of the dataset.
+    is_3D : bool, optional
+        Whether the model is 3D or not, by default False.
+
+    Returns
+    -------
+    dict
+        Model specs compatible with bioimage.io export.
+    """
+    rdf = {
+        "name": "Noise2Void",
+        "description": "Self-supervised denoising.",
+        "license": "BSD-3-Clause",
+        "authors": [
+            {"name": "Alexander Krull"},
+            {"name": "Tim-Oliver Buchholz"},
+            {"name": "Florian Jug"},
+        ],
+        "cite": [
+            {
+                "doi": "10.48550/arXiv.1811.10980",
+                "text": 'A. Krull, T.-O. Buchholz and F. Jug, "Noise2Void - Learning '
+                'Denoising From Single Noisy Images," 2019 IEEE/CVF '
+                "Conference on Computer Vision and Pattern Recognition "
+                "(CVPR), 2019, pp. 2124-2132",
+            }
+        ],
+        # "input_axes": ["bcyx"], <- overriden in save_as_bioimage
+        "preprocessing": [  # for multiple inputs
+            [  # multiple processes per input
+                {
+                    "kwargs": {
+                        "axes": "zyx" if is_3D else "yx",
+                        "mean": [mean],
+                        "mode": "fixed",
+                        "std": [std],
+                    },
+                    "name": "zero_mean_unit_variance",
+                }
+            ]
+        ],
+        # "output_axes": ["bcyx"], <- overriden in save_as_bioimage
+        "postprocessing": [  # for multiple outputs
+            [  # multiple processes per input
+                {
+                    "kwargs": {
+                        "axes": "zyx" if is_3D else "yx",
+                        "gain": [std],
+                        "offset": [mean],
+                    },
+                    "name": "scale_linear",
+                }
+            ]
+        ],
+        "tags": ["unet", "denoising", "Noise2Void", "tensorflow", "napari"],
+    }
+
+    rdf["documentation"] = _get_model_doc(name)
+
+    return rdf

careamics/config/config.py

@@ -13,10 +13,11 @@ from pydantic import (
     model_validator,
 )
 
-from .algorithm import Algorithm
+# ignore typing-only-first-party-import in this file (flake8)
+from .algorithm import Algorithm  # noqa: TCH001
 from .config_filter import paths_to_str
-from .data import Data
-from .training import Training
+from .data import Data  # noqa: TCH001
+from .training import Training  # noqa: TCH001
 
 
 class Configuration(BaseModel):

careamics/config/data.py

@@ -12,7 +12,7 @@ from pydantic import (
     model_validator,
 )
 
-from ..utils import check_axes_validity
+from careamics.utils import check_axes_validity
 
 
 class SupportedExtension(str, Enum):
@@ -23,7 +23,7 @@ class SupportedExtension(str, Enum):
         - tif/tiff: .tiff files.
     """
 
-    TIFF = "tiff"
+    TIFF = "tiff"  # TODO these should be a single one
     TIF = "tif"
 
     @classmethod

careamics/dataset/dataset_utils.py

@@ -46,15 +46,11 @@ def _update_axes(array: np.ndarray, axes: str) -> np.ndarray:
         Updated array.
     """
     # concatenate ST axes to N, return NCZYX
-    if ("S" in axes or "T" in axes) and array.dtype != "O":
+    if "S" in axes or "T" in axes:
         new_axes_len = len(axes.replace("Z", "").replace("YX", ""))
         # TODO test reshape as it can scramble data, moveaxis is probably better
         array = array.reshape(-1, *array.shape[new_axes_len:]).astype(np.float32)
 
-    elif array.dtype == "O":
-        for i in range(len(array)):
-            array[i] = np.expand_dims(array[i], axis=0).astype(np.float32)
-
     else:
         array = np.expand_dims(array, axis=0).astype(np.float32)
 

careamics/dataset/in_memory_dataset.py

@@ -5,8 +5,9 @@ from typing import Callable, Dict, List, Optional, Tuple, Union
 import numpy as np
 import torch
 
-from ..utils import normalize
-from ..utils.logging import get_logger
+from careamics.utils import normalize
+from careamics.utils.logging import get_logger
+
 from .dataset_utils import (
     list_files,
     read_tiff,

careamics/dataset/patching.py

@@ -9,7 +9,8 @@ from typing import Generator, List, Optional, Tuple, Union
 import numpy as np
 from skimage.util import view_as_windows
 
-from ..utils.logging import get_logger
+from careamics.utils.logging import get_logger
+
 from .extraction_strategy import ExtractionStrategy
 
 logger = get_logger(__name__)
@@ -481,13 +482,11 @@ def generate_patches(
         elif patch_extraction_method == ExtractionStrategy.SEQUENTIAL:
             patches = _extract_patches_sequential(sample, patch_size=patch_size)
 
-        elif patch_extraction_method == ExtractionStrategy.RANDOM:
+        else:
+            # random patching
             patches = _extract_patches_random(sample, patch_size=patch_size)
 
-        if patches is None:
-            raise ValueError("No patch generated")
-
         return patches
     else:
-        # no patching
+        # no patching, return a generator for the sample
         return (sample for _ in range(1))

careamics/dataset/prepare_dataset.py

@@ -6,9 +6,10 @@ Methods to set up the datasets for training, validation and prediction.
 from pathlib import Path
 from typing import List, Optional, Union
 
-from ..config import Configuration
-from ..manipulation import default_manipulate
-from ..utils import check_tiling_validity
+from careamics.config import Configuration
+from careamics.manipulation import default_manipulate
+from careamics.utils import check_tiling_validity
+
 from .extraction_strategy import ExtractionStrategy
 from .in_memory_dataset import InMemoryDataset
 from .tiff_dataset import TiffDataset

careamics/dataset/tiff_dataset.py

@@ -10,8 +10,9 @@ from typing import Callable, Dict, Generator, List, Optional, Tuple, Union
 import numpy as np
 import torch
 
-from ..utils import normalize
-from ..utils.logging import get_logger
+from careamics.utils import normalize
+from careamics.utils.logging import get_logger
+
 from .dataset_utils import (
     list_files,
     read_tiff,
@@ -53,7 +54,7 @@ class TiffDataset(torch.utils.data.IterableDataset):
     def __init__(
         self,
         data_path: Union[str, Path],
-        data_format: str,
+        data_format: str,  # TODO: TiffDataset should not know that they are tiff
         axes: str,
         patch_extraction_method: Union[ExtractionStrategy, None],
         patch_size: Optional[Union[List[int], Tuple[int]]] = None,