careamics 0.0.14__py3-none-any.whl → 0.0.16__py3-none-any.whl

careamics/careamist.py

@@ -2,7 +2,7 @@
 
 from collections.abc import Callable
 from pathlib import Path
-from typing import Any, Literal, Optional, Union, overload
+from typing import Any, Literal, Union, overload
 
 import numpy as np
 from numpy.typing import NDArray
@@ -79,8 +79,8 @@ class CAREamist:
     def __init__(  # numpydoc ignore=GL08
         self,
         source: Union[Path, str],
-        work_dir: Optional[Union[Path, str]] = None,
-        callbacks: Optional[list[Callback]] = None,
+        work_dir: Union[Path, str] | None = None,
+        callbacks: list[Callback] | None = None,
         enable_progress_bar: bool = True,
     ) -> None: ...
 
@@ -88,16 +88,16 @@ class CAREamist:
     def __init__(  # numpydoc ignore=GL08
         self,
         source: Configuration,
-        work_dir: Optional[Union[Path, str]] = None,
-        callbacks: Optional[list[Callback]] = None,
+        work_dir: Union[Path, str] | None = None,
+        callbacks: list[Callback] | None = None,
         enable_progress_bar: bool = True,
     ) -> None: ...
 
     def __init__(
         self,
         source: Union[Path, str, Configuration],
-        work_dir: Optional[Union[Path, str]] = None,
-        callbacks: Optional[list[Callback]] = None,
+        work_dir: Union[Path, str] | None = None,
+        callbacks: list[Callback] | None = None,
         enable_progress_bar: bool = True,
     ) -> None:
         """
@@ -208,25 +208,19 @@ class CAREamist:
 
         # instantiate trainer
         self.trainer = Trainer(
-            max_epochs=self.cfg.training_config.num_epochs,
-            precision=self.cfg.training_config.precision,
-            max_steps=self.cfg.training_config.max_steps,
-            check_val_every_n_epoch=self.cfg.training_config.check_val_every_n_epoch,
             enable_progress_bar=enable_progress_bar,
-            accumulate_grad_batches=self.cfg.training_config.accumulate_grad_batches,
-            gradient_clip_val=self.cfg.training_config.gradient_clip_val,
-            gradient_clip_algorithm=self.cfg.training_config.gradient_clip_algorithm,
             callbacks=self.callbacks,
             default_root_dir=self.work_dir,
             logger=experiment_logger,
+            **self.cfg.training_config.lightning_trainer_config or {},
         )
 
         # place holder for the datamodules
-        self.train_datamodule: Optional[TrainDataModule] = None
-        self.pred_datamodule: Optional[PredictDataModule] = None
+        self.train_datamodule: TrainDataModule | None = None
+        self.pred_datamodule: PredictDataModule | None = None
 
     def _define_callbacks(
-        self, callbacks: Optional[list[Callback]], enable_progress_bar: bool
+        self, callbacks: list[Callback] | None, enable_progress_bar: bool
     ) -> None:
         """Define the callbacks for the training loop.
 
@@ -264,7 +258,7 @@ class CAREamist:
                 HyperParametersCallback(self.cfg),
                 ModelCheckpoint(
                     dirpath=self.work_dir / Path("checkpoints"),
-                    filename=self.cfg.experiment_name,
+                    filename=f"{self.cfg.experiment_name}_{{epoch:02d}}_step_{{step}}",
                     **self.cfg.training_config.checkpoint_callback.model_dump(),
                 ),
             ]
@@ -288,11 +282,11 @@ class CAREamist:
     def train(
         self,
         *,
-        datamodule: Optional[TrainDataModule] = None,
-        train_source: Optional[Union[Path, str, NDArray]] = None,
-        val_source: Optional[Union[Path, str, NDArray]] = None,
-        train_target: Optional[Union[Path, str, NDArray]] = None,
-        val_target: Optional[Union[Path, str, NDArray]] = None,
+        datamodule: TrainDataModule | None = None,
+        train_source: Union[Path, str, NDArray] | None = None,
+        val_source: Union[Path, str, NDArray] | None = None,
+        train_target: Union[Path, str, NDArray] | None = None,
+        val_target: Union[Path, str, NDArray] | None = None,
         use_in_memory: bool = True,
         val_percentage: float = 0.1,
         val_minimum_split: int = 1,
@@ -443,9 +437,9 @@ class CAREamist:
     def _train_on_array(
         self,
         train_data: NDArray,
-        val_data: Optional[NDArray] = None,
-        train_target: Optional[NDArray] = None,
-        val_target: Optional[NDArray] = None,
+        val_data: NDArray | None = None,
+        train_target: NDArray | None = None,
+        val_target: NDArray | None = None,
         val_percentage: float = 0.1,
         val_minimum_split: int = 5,
     ) -> None:
@@ -484,9 +478,9 @@ class CAREamist:
     def _train_on_path(
         self,
         path_to_train_data: Union[Path, str],
-        path_to_val_data: Optional[Union[Path, str]] = None,
-        path_to_train_target: Optional[Union[Path, str]] = None,
-        path_to_val_target: Optional[Union[Path, str]] = None,
+        path_to_val_data: Union[Path, str] | None = None,
+        path_to_train_target: Union[Path, str] | None = None,
+        path_to_val_target: Union[Path, str] | None = None,
         use_in_memory: bool = True,
         val_percentage: float = 0.1,
         val_minimum_split: int = 1,
@@ -549,13 +543,13 @@ class CAREamist:
         source: Union[Path, str],
         *,
         batch_size: int = 1,
-        tile_size: Optional[tuple[int, ...]] = None,
-        tile_overlap: Optional[tuple[int, ...]] = (48, 48),
-        axes: Optional[str] = None,
-        data_type: Optional[Literal["tiff", "custom"]] = None,
+        tile_size: tuple[int, ...] | None = None,
+        tile_overlap: tuple[int, ...] | None = (48, 48),
+        axes: str | None = None,
+        data_type: Literal["tiff", "custom"] | None = None,
         tta_transforms: bool = False,
-        dataloader_params: Optional[dict] = None,
-        read_source_func: Optional[Callable] = None,
+        dataloader_params: dict | None = None,
+        read_source_func: Callable | None = None,
         extension_filter: str = "",
     ) -> Union[list[NDArray], NDArray]: ...
 
@@ -565,12 +559,12 @@ class CAREamist:
         source: NDArray,
         *,
         batch_size: int = 1,
-        tile_size: Optional[tuple[int, ...]] = None,
-        tile_overlap: Optional[tuple[int, ...]] = (48, 48),
-        axes: Optional[str] = None,
-        data_type: Optional[Literal["array"]] = None,
+        tile_size: tuple[int, ...] | None = None,
+        tile_overlap: tuple[int, ...] | None = (48, 48),
+        axes: str | None = None,
+        data_type: Literal["array"] | None = None,
         tta_transforms: bool = False,
-        dataloader_params: Optional[dict] = None,
+        dataloader_params: dict | None = None,
     ) -> Union[list[NDArray], NDArray]: ...
 
     def predict(
@@ -578,13 +572,13 @@ class CAREamist:
         source: Union[PredictDataModule, Path, str, NDArray],
         *,
         batch_size: int = 1,
-        tile_size: Optional[tuple[int, ...]] = None,
-        tile_overlap: Optional[tuple[int, ...]] = (48, 48),
-        axes: Optional[str] = None,
-        data_type: Optional[Literal["array", "tiff", "custom"]] = None,
+        tile_size: tuple[int, ...] | None = None,
+        tile_overlap: tuple[int, ...] | None = (48, 48),
+        axes: str | None = None,
+        data_type: Literal["array", "tiff", "custom"] | None = None,
         tta_transforms: bool = False,
-        dataloader_params: Optional[dict] = None,
-        read_source_func: Optional[Callable] = None,
+        dataloader_params: dict | None = None,
+        read_source_func: Callable | None = None,
         extension_filter: str = "",
         **kwargs: Any,
     ) -> Union[list[NDArray], NDArray]:
@@ -704,18 +698,18 @@ class CAREamist:
         source: Union[PredictDataModule, Path, str],
         *,
         batch_size: int = 1,
-        tile_size: Optional[tuple[int, ...]] = None,
-        tile_overlap: Optional[tuple[int, ...]] = (48, 48),
-        axes: Optional[str] = None,
-        data_type: Optional[Literal["tiff", "custom"]] = None,
+        tile_size: tuple[int, ...] | None = None,
+        tile_overlap: tuple[int, ...] | None = (48, 48),
+        axes: str | None = None,
+        data_type: Literal["tiff", "custom"] | None = None,
         tta_transforms: bool = False,
-        dataloader_params: Optional[dict] = None,
-        read_source_func: Optional[Callable] = None,
+        dataloader_params: dict | None = None,
+        read_source_func: Callable | None = None,
         extension_filter: str = "",
         write_type: Literal["tiff", "custom"] = "tiff",
-        write_extension: Optional[str] = None,
-        write_func: Optional[WriteFunc] = None,
-        write_func_kwargs: Optional[dict[str, Any]] = None,
+        write_extension: str | None = None,
+        write_func: WriteFunc | None = None,
+        write_func_kwargs: dict[str, Any] | None = None,
         prediction_dir: Union[Path, str] = "predictions",
         **kwargs,
     ) -> None:
@@ -823,12 +817,12 @@ class CAREamist:
 
         # extract file names
         source_path: Union[Path, str, NDArray]
-        source_data_type: Literal["array", "tiff", "custom"]
+        source_data_type: Literal["array", "tiff", "czi", "custom"]
         if isinstance(source, PredictDataModule):
             source_path = source.pred_data
             source_data_type = source.data_type
             extension_filter = source.extension_filter
-        elif isinstance(source, str | Path):
+        elif isinstance(source, (str | Path)):
             source_path = source
             source_data_type = data_type or self.cfg.data_config.data_type
             extension_filter = SupportedData.get_extension_pattern(
@@ -841,7 +835,7 @@ class CAREamist:
             raise ValueError(
                 "Predicting to disk is not supported for input type 'array'."
             )
-        assert isinstance(source_path, str | Path)  # because data_type != "array"
+        assert isinstance(source_path, (Path | str))  # because data_type != "array"
         source_path = Path(source_path)
 
         file_paths = list_files(source_path, source_data_type, extension_filter)
@@ -879,14 +873,14 @@ class CAREamist:
 
     def export_to_bmz(
         self,
-        path_to_archive: Union[Path, str],
+        path_to_archive: Union[Path | str],
         friendly_model_name: str,
         input_array: NDArray,
         authors: list[dict],
         general_description: str,
         data_description: str,
-        covers: Optional[list[Union[Path, str]]] = None,
-        channel_names: Optional[list[str]] = None,
+        covers: list[Union[Path, str]] | None = None,
+        channel_names: list[str] | None = None,
         model_version: str = "0.1.0",
     ) -> None:
         """Export the model to the BioImage Model Zoo format.

careamics/cli/conf.py

@@ -3,7 +3,7 @@
 import sys
 from dataclasses import dataclass
 from pathlib import Path
-from typing import Annotated, Optional
+from typing import Annotated
 
 import click
 import typer
@@ -116,7 +116,11 @@ def care(  # numpydoc ignore=PR01
         ),
     ],
     batch_size: Annotated[int, typer.Option(help="Batch size.")],
-    num_epochs: Annotated[int, typer.Option(help="Number of epochs.")],
+    num_epochs: Annotated[int, typer.Option(help="Number of epochs.")] = 100,
+    num_steps: Annotated[
+        int | None,
+        typer.Option(help="Number of batches per epoch (limit_train_batches)."),
+    ] = None,
     data_type: Annotated[
         click.Choice,
         typer.Option(click_type=click.Choice(["tiff"]), help="Type of the data."),
@@ -135,10 +139,10 @@ def care(  # numpydoc ignore=PR01
         ),
     ] = "mae",
     n_channels_in: Annotated[
-        Optional[int], typer.Option(help="Number of channels in")
+        int | None, typer.Option(help="Number of channels in")
     ] = None,
     n_channels_out: Annotated[
-        Optional[int], typer.Option(help="Number of channels out")
+        int | None, typer.Option(help="Number of channels out")
     ] = None,
     logger: Annotated[
         click.Choice,
@@ -175,6 +179,7 @@ def care(  # numpydoc ignore=PR01
         patch_size=patch_size,
         batch_size=batch_size,
         num_epochs=num_epochs,
+        num_steps=num_steps,
         # TODO: fix choosing augmentations
         augmentations=None if use_augmentations else [],
         independent_channels=independent_channels,
@@ -203,7 +208,11 @@ def n2n(  # numpydoc ignore=PR01
         ),
     ],
     batch_size: Annotated[int, typer.Option(help="Batch size.")],
-    num_epochs: Annotated[int, typer.Option(help="Number of epochs.")],
+    num_epochs: Annotated[int, typer.Option(help="Number of epochs.")] = 100,
+    num_steps: Annotated[
+        int | None,
+        typer.Option(help="Number of batches per epoch (limit_train_batches)."),
+    ] = None,
     data_type: Annotated[
         click.Choice,
         typer.Option(click_type=click.Choice(["tiff"]), help="Type of the data."),
@@ -222,10 +231,10 @@ def n2n(  # numpydoc ignore=PR01
         ),
     ] = "mae",
     n_channels_in: Annotated[
-        Optional[int], typer.Option(help="Number of channels in")
+        int | None, typer.Option(help="Number of channels in")
     ] = None,
     n_channels_out: Annotated[
-        Optional[int], typer.Option(help="Number of channels out")
+        int | None, typer.Option(help="Number of channels out")
     ] = None,
     logger: Annotated[
         click.Choice,
@@ -259,6 +268,7 @@ def n2n(  # numpydoc ignore=PR01
         patch_size=patch_size,
         batch_size=batch_size,
         num_epochs=num_epochs,
+        num_steps=num_steps,
         # TODO: fix choosing augmentations
         augmentations=None if use_augmentations else [],
         independent_channels=independent_channels,
@@ -287,7 +297,11 @@ def n2v(  # numpydoc ignore=PR01
         ),
     ],
     batch_size: Annotated[int, typer.Option(help="Batch size.")],
-    num_epochs: Annotated[int, typer.Option(help="Number of epochs.")],
+    num_epochs: Annotated[int, typer.Option(help="Number of epochs.")] = 100,
+    num_steps: Annotated[
+        int | None,
+        typer.Option(help="Number of batches per epoch (limit_train_batches)."),
+    ] = None,
     data_type: Annotated[
         click.Choice,
         typer.Option(click_type=click.Choice(["tiff"]), help="Type of the data."),
@@ -300,7 +314,7 @@ def n2v(  # numpydoc ignore=PR01
     ] = True,
     use_n2v2: Annotated[bool, typer.Option(help="Whether to use N2V2")] = False,
     n_channels: Annotated[
-        Optional[int], typer.Option(help="Number of channels (in and out)")
+        int | None, typer.Option(help="Number of channels (in and out)")
     ] = None,
     roi_size: Annotated[int, typer.Option(help="N2V pixel manipulation area.")] = 11,
     masked_pixel_percentage: Annotated[
@@ -364,6 +378,7 @@ def n2v(  # numpydoc ignore=PR01
         patch_size=patch_size,
         batch_size=batch_size,
         num_epochs=num_epochs,
+        num_steps=num_steps,
         # TODO: fix choosing augmentations
         augmentations=None if use_augmentations else [],
         independent_channels=independent_channels,

careamics/cli/main.py

@@ -7,7 +7,7 @@ its implementation is contained in the conf.py file.
 """
 
 from pathlib import Path
-from typing import Annotated, Optional
+from typing import Annotated
 
 import click
 import typer
@@ -47,7 +47,7 @@ def train(  # numpydoc ignore=PR01
         ),
     ],
     train_target: Annotated[
-        Optional[Path],
+        Path | None,
         typer.Option(
             "--train-target",
             "-tt",
@@ -58,7 +58,7 @@ def train(  # numpydoc ignore=PR01
         ),
     ] = None,
     val_source: Annotated[
-        Optional[Path],
+        Path | None,
         typer.Option(
             "--val-source",
             "-vs",
@@ -69,7 +69,7 @@ def train(  # numpydoc ignore=PR01
         ),
     ] = None,
     val_target: Annotated[
-        Optional[Path],
+        Path | None,
         typer.Option(
             "--val-target",
             "-vt",
@@ -96,7 +96,7 @@ def train(  # numpydoc ignore=PR01
         typer.Option(help="Minimum number of files to use for validation,"),
     ] = 1,
     work_dir: Annotated[
-        Optional[Path],
+        Path | None,
         typer.Option(
             "--work-dir",
             "-wd",
@@ -142,7 +142,7 @@ def predict(  # numpydoc ignore=PR01
     ],
     batch_size: Annotated[int, typer.Option(help="Batch size.")] = 1,
     tile_size: Annotated[
-        Optional[click.Tuple],
+        click.Tuple | None,
         typer.Option(
             help=(
                 "Size of the tiles to use for prediction, (if the data "
@@ -164,7 +164,7 @@ def predict(  # numpydoc ignore=PR01
         ),
     ] = (48, 48, -1),
     axes: Annotated[
-        Optional[str],
+        str | None,
         typer.Option(
             help="Axes of the input data. If unused the data is assumed to have the "
             "same axes as the original training data."
@@ -190,7 +190,7 @@ def predict(  # numpydoc ignore=PR01
     ] = "tiff",
     # TODO: could make dataloader_params as json, necessary?
     work_dir: Annotated[
-        Optional[Path],
+        Path | None,
         typer.Option(
             "--work-dir",
             "-wd",

careamics/cli/utils.py

@@ -1,11 +1,9 @@
 """Utility functions for the CAREamics CLI."""
 
-from typing import Optional
-
 
 def handle_2D_3D_callback(
-    value: Optional[tuple[int, int, int]],
-) -> Optional[tuple[int, ...]]:
+    value: tuple[int, int, int] | None,
+) -> tuple[int, ...] | None:
     """
     Callback for options that require 2D or 3D inputs.
 

careamics/config/__init__.py

@@ -12,8 +12,10 @@ __all__ = [
     "Configuration",
     "DataConfig",
     "GaussianMixtureNMConfig",
+    "HDNAlgorithm",
     "InferenceConfig",
     "LVAELossConfig",
+    "MicroSplitAlgorithm",
     "MultiChannelNMConfig",
     "N2NAlgorithm",
     "N2VAlgorithm",
@@ -22,6 +24,8 @@ __all__ = [
     "VAEBasedAlgorithm",
     "algorithm_factory",
     "create_care_configuration",
+    "create_hdn_configuration",
+    "create_microsplit_configuration",
     "create_n2n_configuration",
     "create_n2v_configuration",
     "load_configuration",
@@ -30,6 +34,8 @@ __all__ = [
 
 from .algorithms import (
     CAREAlgorithm,
+    HDNAlgorithm,
+    MicroSplitAlgorithm,
     N2NAlgorithm,
     N2VAlgorithm,
     UNetBasedAlgorithm,
@@ -40,6 +46,8 @@ from .configuration import Configuration
 from .configuration_factories import (
     algorithm_factory,
     create_care_configuration,
+    create_hdn_configuration,
+    create_microsplit_configuration,
     create_n2n_configuration,
     create_n2v_configuration,
 )

careamics/config/algorithms/__init__.py

@@ -2,6 +2,8 @@
 
 __all__ = [
     "CAREAlgorithm",
+    "HDNAlgorithm",
+    "MicroSplitAlgorithm",
     "N2NAlgorithm",
     "N2VAlgorithm",
     "UNetBasedAlgorithm",
@@ -9,6 +11,8 @@ __all__ = [
 ]
 
 from .care_algorithm_model import CAREAlgorithm
+from .hdn_algorithm_model import HDNAlgorithm
+from .microsplit_algorithm_model import MicroSplitAlgorithm
 from .n2n_algorithm_model import N2NAlgorithm
 from .n2v_algorithm_model import N2VAlgorithm
 from .unet_algorithm_model import UNetBasedAlgorithm

careamics/config/algorithms/hdn_algorithm_model.py

@@ -0,0 +1,103 @@
+"""HDN algorithm configuration."""
+
+from typing import Literal
+
+from bioimageio.spec.generic.v0_3 import CiteEntry
+from pydantic import ConfigDict
+
+from careamics.config.algorithms.vae_algorithm_model import VAEBasedAlgorithm
+from careamics.config.architectures import LVAEModel
+from careamics.config.loss_model import LVAELossConfig
+
+HDN = "Hierarchical DivNoising"
+
+HDN_DESCRIPTION = (
+    "HDN leverages a hierarchical VAE to perform image "
+    "restoration. It is designed to be interpretable and unsupervised, "
+    "making it suitable for a wide range of microscopy images."
+)
+HDN_REF = CiteEntry(
+    text='Prakash, M., Delbracio, M., Milanfar, P., Jug, F. 2022. "Interpretable '
+    'Unsupervised Diversity Denoising and Artefact Removal." The International '
+    "Conference on Learning Representations (ICLR).",
+    doi="10.1561/2200000056",
+)
+
+
+class HDNAlgorithm(VAEBasedAlgorithm):
+    """HDN algorithm configuration."""
+
+    model_config = ConfigDict(validate_assignment=True)
+
+    algorithm: Literal["hdn"] = "hdn"
+
+    loss: LVAELossConfig
+
+    model: LVAEModel  # TODO add validators
+
+    is_supervised: bool = False
+
+    def get_algorithm_friendly_name(self) -> str:
+        """
+        Get the algorithm friendly name.
+
+        Returns
+        -------
+        str
+            Friendly name of the algorithm.
+        """
+        return HDN
+
+    def get_algorithm_keywords(self) -> list[str]:
+        """
+        Get algorithm keywords.
+
+        Returns
+        -------
+        list[str]
+            List of keywords.
+        """
+        return [
+            "restoration",
+            "VAE",
+            "3D" if self.model.is_3D() else "2D",
+            "CAREamics",
+            "pytorch",
+        ]
+
+    def get_algorithm_references(self) -> str:
+        """
+        Get the algorithm references.
+
+        This is used to generate the README of the BioImage Model Zoo export.
+
+        Returns
+        -------
+        str
+            Algorithm references.
+        """
+        return HDN_REF.text + " doi: " + HDN_REF.doi
+
+    def get_algorithm_citations(self) -> list[CiteEntry]:
+        """
+        Return a list of citation entries of the current algorithm.
+
+        This is used to generate the model description for the BioImage Model Zoo.
+
+        Returns
+        -------
+        List[CiteEntry]
+            List of citation entries.
+        """
+        return [HDN_REF]
+
+    def get_algorithm_description(self) -> str:
+        """
+        Get the algorithm description.
+
+        Returns
+        -------
+        str
+            Algorithm description.
+        """
+        return HDN_DESCRIPTION