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

careamics/config/data/data_model.py

@@ -6,7 +6,7 @@ import os
 import sys
 from collections.abc import Sequence
 from pprint import pformat
-from typing import Annotated, Any, Literal, Union
+from typing import Annotated, Any, Literal, Self, Union
 from warnings import warn
 
 import numpy as np
@@ -19,7 +19,6 @@ from pydantic import (
     field_validator,
     model_validator,
 )
-from typing_extensions import Self
 
 from ..transformations import XYFlipModel, XYRandomRotate90Model
 from ..validators import check_axes_validity, patch_size_ge_than_8_power_of_2

careamics/config/data/ng_data_model.py

@@ -4,7 +4,7 @@ from __future__ import annotations
 
 from collections.abc import Sequence
 from pprint import pformat
-from typing import Annotated, Any, Literal, Union
+from typing import Annotated, Any, Literal, Self, Union
 from warnings import warn
 
 import numpy as np
@@ -17,7 +17,6 @@ from pydantic import (
     field_validator,
     model_validator,
 )
-from typing_extensions import Self
 
 from ..transformations import XYFlipModel, XYRandomRotate90Model
 from ..validators import check_axes_validity

careamics/config/inference_model.py

@@ -2,10 +2,9 @@
 
 from __future__ import annotations
 
-from typing import Any, Literal, Union
+from typing import Any, Literal, Self, Union
 
 from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator
-from typing_extensions import Self
 
 from .validators import check_axes_validity, patch_size_ge_than_8_power_of_2
 

careamics/config/likelihood_model.py

@@ -50,11 +50,11 @@ class NMLikelihoodConfig(BaseModel):
     model_config = ConfigDict(validate_assignment=True, arbitrary_types_allowed=True)
 
     # TODO remove and use as parameters to the likelihood functions?
-    data_mean: Tensor = torch.zeros(1)
+    data_mean: Tensor | None = None
     """The mean of the data, used to unnormalize data for noise model evaluation.
     Shape is (target_ch,) (or (1, target_ch, [1], 1, 1))."""
 
     # TODO remove and use as parameters to the likelihood functions?
-    data_std: Tensor = torch.ones(1)
+    data_std: Tensor | None = None
     """The standard deviation of the data, used to unnormalize data for noise
     model evaluation. Shape is (target_ch,) (or (1, target_ch, [1], 1, 1))."""

careamics/config/loss_model.py

@@ -35,7 +35,9 @@ class LVAELossConfig(BaseModel):
         validate_assignment=True, validate_default=True, arbitrary_types_allowed=True
     )
 
-    loss_type: Literal["musplit", "denoisplit", "denoisplit_musplit"]
+    loss_type: Literal[
+        "hdn", "microsplit", "musplit", "denoisplit", "denoisplit_musplit"
+    ]
     """Type of loss to use for LVAE."""
 
     reconstruction_weight: float = 1.0
@@ -50,7 +52,9 @@ class LVAELossConfig(BaseModel):
     """Weight for the denoiSplit loss (used in the muSplit-deonoiSplit loss)."""
     kl_params: KLLossConfig = KLLossConfig()
     """KL loss configuration."""
-
+    # TODO revisit weights for the losses
     # TODO: remove?
     non_stochastic: bool = False
     """Whether to sample latents and compute KL."""
+
+    # TODO what are the correct parameters for HDN ?

careamics/config/nm_model.py

@@ -1,7 +1,7 @@
 """Noise models config."""
 
 from pathlib import Path
-from typing import Annotated, Literal, Union
+from typing import Annotated, Literal, Self, Union
 
 import numpy as np
 import torch
@@ -11,6 +11,7 @@ from pydantic import (
     Field,
     PlainSerializer,
     PlainValidator,
+    model_validator,
 )
 
 from careamics.utils.serializers import _array_to_json, _to_numpy
@@ -86,6 +87,30 @@ class GaussianMixtureNMConfig(BaseModel):
     tol: float = Field(default=1e-10)
     """Tolerance used in the computation of the noise model likelihood."""
 
+    @model_validator(mode="after")
+    def validate_path(self: Self) -> Self:
+        """Validate that the path points to a valid .npz file if provided.
+
+        Returns
+        -------
+        Self
+            Returns itself.
+
+        Raises
+        ------
+        ValueError
+            If the path is provided but does not point to a valid .npz file.
+        """
+        if self.path is not None:
+            path = Path(self.path)
+            if not path.exists():
+                raise ValueError(f"Path {path} does not exist.")
+            if path.suffix != ".npz":
+                raise ValueError(f"Path {path} must point to a .npz file.")
+            if not path.is_file():
+                raise ValueError(f"Path {path} must point to a file.")
+        return self
+
     # @model_validator(mode="after")
     # def validate_path_to_pretrained_vs_training_data(self: Self) -> Self:
     #     """Validate paths provided in the config.

careamics/config/optimizer_models.py

@@ -2,7 +2,7 @@
 
 from __future__ import annotations
 
-from typing import Literal
+from typing import Literal, Self
 
 from pydantic import (
     BaseModel,
@@ -13,7 +13,6 @@ from pydantic import (
     model_validator,
 )
 from torch import optim
-from typing_extensions import Self
 
 from careamics.utils.torch_utils import filter_parameters
 

careamics/config/support/supported_algorithms.py

@@ -26,9 +26,11 @@ class SupportedAlgorithm(str, BaseEnum):
     MUSPLIT = "musplit"
     """An image splitting approach based on ladder VAE architectures."""
 
+    MICROSPLIT = "microsplit"
+    """A micro-level image splitting approach based on ladder VAE architectures."""
+
     DENOISPLIT = "denoisplit"
     """An image splitting and denoising approach based on ladder VAE architectures."""
 
-    # PN2V = "pn2v"
-    # HDN = "hdn"
-    # SEG = "segmentation"
+    HDN = "hdn"
+    """Hierarchical Denoising Network, an unsupervised denoising algorithm"""

careamics/config/support/supported_losses.py

@@ -21,9 +21,12 @@ class SupportedLoss(str, BaseEnum):
     MAE = "mae"
     N2V = "n2v"
     # PN2V = "pn2v"
-    # HDN = "hdn"
+    HDN = "hdn"
     MUSPLIT = "musplit"
+    MICROSPLIT = "microsplit"
     DENOISPLIT = "denoisplit"
-    DENOISPLIT_MUSPLIT = "denoisplit_musplit"
+    DENOISPLIT_MUSPLIT = (
+        "denoisplit_musplit"  # TODO refac losses, leave only microsplit
+    )
     # CE = "ce"
     # DICE = "dice"

careamics/config/training_model.py

@@ -3,9 +3,9 @@
 from __future__ import annotations
 
 from pprint import pformat
-from typing import Literal, Union
+from typing import Literal
 
-from pydantic import BaseModel, ConfigDict, Field, field_validator
+from pydantic import BaseModel, ConfigDict, Field
 
 from .callback_model import CheckpointModel, EarlyStoppingModel
 
@@ -29,26 +29,15 @@ class TrainingConfig(BaseModel):
     model_config = ConfigDict(
         validate_assignment=True,
     )
+    lightning_trainer_config: dict | None = None
+    """Configuration for the PyTorch Lightning Trainer, following PyTorch Lightning
+    Trainer class"""
 
-    num_epochs: int = Field(default=20, ge=1)
-    """Number of epochs, greater than 0."""
-
-    precision: Literal["64", "32", "16-mixed", "bf16-mixed"] = Field(default="32")
-    """Numerical precision"""
-    max_steps: int = Field(default=-1, ge=-1)
-    """Maximum number of steps to train for. -1 means no limit."""
-    check_val_every_n_epoch: int = Field(default=1, ge=1)
-    """Validation step frequency."""
-    accumulate_grad_batches: int = Field(default=1, ge=1)
-    """Number of batches to accumulate gradients over before stepping the optimizer."""
-    gradient_clip_val: Union[int, float] | None = None
-    """The value to which to clip the gradient"""
-    gradient_clip_algorithm: Literal["value", "norm"] = "norm"
-    """The algorithm to use for gradient clipping (see lightning `Trainer`)."""
     logger: Literal["wandb", "tensorboard"] | None = None
     """Logger to use during training. If None, no logger will be used. Available
     loggers are defined in SupportedLogger."""
 
+    # Only basic callbacks
     checkpoint_callback: CheckpointModel = CheckpointModel()
     """Checkpoint callback configuration, following PyTorch Lightning Checkpoint
     callback."""
@@ -78,22 +67,3 @@ class TrainingConfig(BaseModel):
             Whether the logger is defined or not.
         """
         return self.logger is not None
-
-    @field_validator("max_steps")
-    @classmethod
-    def validate_max_steps(cls, max_steps: int) -> int:
-        """Validate the max_steps parameter.
-
-        Parameters
-        ----------
-        max_steps : int
-            Maximum number of steps to train for. -1 means no limit.
-
-        Returns
-        -------
-        int
-            Validated max_steps.
-        """
-        if max_steps == 0:
-            raise ValueError("max_steps must be greater than 0. Use -1 for no limit.")
-        return max_steps

careamics/config/transformations/normalize_model.py

@@ -1,9 +1,8 @@
 """Pydantic model for the Normalize transform."""
 
-from typing import Literal
+from typing import Literal, Self
 
 from pydantic import ConfigDict, Field, model_validator
-from typing_extensions import Self
 
 from .transform_model import TransformModel
 

careamics/dataset_ng/patch_extractor/demo_custom_image_stack_loader.py

@@ -7,7 +7,7 @@ import matplotlib.pyplot as plt
 import numpy as np
 import zarr
 from numpy.typing import NDArray
-from zarr.storage import FSStore
+from zarr.storage import FsspecStore
 
 from careamics.config import DataConfig
 from careamics.config.support import SupportedData
@@ -20,7 +20,7 @@ from careamics.dataset_ng.patch_extractor.patch_extractor_factory import (
 
 # %%
 def create_zarr_array(file_path: Path, data_path: str, data: NDArray):
-    store = FSStore(url=file_path.resolve())
+    store = FsspecStore.from_url(url=file_path.resolve())
     # create array
     array = zarr.create(
         store=store,
@@ -61,7 +61,7 @@ if not file_path.is_file() and not file_path.is_dir():
 # ### Make sure file exists
 
 # %%
-store = FSStore(url=file_path.resolve(), mode="r")
+store = FsspecStore.from_url(url=file_path.resolve(), mode="r")
 
 # %%
 list(store.keys())
@@ -72,7 +72,7 @@ list(store.keys())
 
 # %%
 class ZarrSource(TypedDict):
-    store: FSStore
+    store: FsspecStore
     data_paths: Sequence[str]
 
 

careamics/dataset_ng/patch_extractor/image_stack/in_memory_image_stack.py

@@ -1,9 +1,8 @@
 from collections.abc import Sequence
 from pathlib import Path
-from typing import Any, Literal, Union
+from typing import Any, Literal, Self, Union
 
 from numpy.typing import DTypeLike, NDArray
-from typing_extensions import Self
 
 from careamics.dataset.dataset_utils import reshape_array
 from careamics.file_io.read import ReadFunc, read_tiff

careamics/dataset_ng/patch_extractor/image_stack/zarr_image_stack.py

@@ -1,11 +1,11 @@
 from collections.abc import Sequence
 from pathlib import Path
-from typing import Union
+from typing import Self, Union
 
+import validators
 import zarr
-import zarr.storage
 from numpy.typing import NDArray
-from typing_extensions import Self
+from zarr.storage import FsspecStore, LocalStore
 
 from careamics.dataset.dataset_utils import reshape_array
 
@@ -15,9 +15,10 @@ class ZarrImageStack:
     A class for extracting patches from an image stack that is stored as a zarr array.
     """
 
-    # TODO: keeping store type narrow so that it has the path attribute
-    #   base zarr store is zarr.storage.Store, includes MemoryStore
-    def __init__(self, store: zarr.storage.FSStore, data_path: str, axes: str):
+    # TODO: We should keep store type narrow
+    #   - in zarr v3, does zarr.storage.Store exists and has the path attribute?
+    #   - can we declare a narrow type rather than a union?
+    def __init__(self, store: LocalStore | FsspecStore, data_path: str, axes: str):
         self._store = store
         self._array = zarr.open_array(store=self._store, path=data_path, mode="r")
         # TODO: validate axes
@@ -46,8 +47,33 @@ class ZarrImageStack:
         Assumes the path only contains 1 image.
 
         Path can be to a local file, or it can be a URL to a zarr stored in the cloud.
+
+        Parameters
+        ----------
+        path : Union[Path, str]
+            Path to the root of the OME-Zarr, local file or url.
+
+        Returns
+        -------
+        ZarrImageStack
+            Initialised ZarrImageStack.
+
+        Raises
+        ------
+        ValueError
+            If the path does not exist or is not a valid URL.
+        ValueError
+            If the OME-Zarr at the path does not contain the attribute 'multiscales'.
         """
-        store = zarr.storage.FSStore(url=path)
+        if Path(path).is_file():
+            store = zarr.storage.LocalStore(root=Path(path).resolve())
+        elif validators.url(path):
+            store = zarr.storage.FsspecStore.from_url(url=path)
+        else:
+            raise ValueError(
+                f"Path '{path}' is neither an existing file nor a valid URL."
+            )
+
         group = zarr.open_group(store=store, mode="r")
         if "multiscales" not in group.attrs:
             raise ValueError(

careamics/dataset_ng/patch_extractor/image_stack_loader.py

@@ -38,7 +38,7 @@ class ImageStackLoader(Protocol[P, GenericImageStack]):
 
     >>> from typing import TypedDict
 
-    >>> from zarr.storage import FSStore
+    >>> from zarr.storage import FsspecStore
 
     >>> from careamics.config import DataConfig
     >>> from careamics.dataset_ng.patch_extractor.image_stack import ZarrImageStack
@@ -46,7 +46,7 @@ class ImageStackLoader(Protocol[P, GenericImageStack]):
     >>> # Define a zarr source
     >>> # It encompasses multiple arguments that determine what data will be loaded
     >>> class ZarrSource(TypedDict):
-    ...     store: FSStore
+    ...     store: FsspecStore
     ...     data_paths: Sequence[str]
 
     >>> def custom_image_stack_loader(

careamics/file_io/read/__init__.py

@@ -9,4 +9,3 @@ __all__ = [
 
 from .get_func import ReadFunc, get_read_func
 from .tiff import read_tiff
-from .zarr import read_zarr

careamics/lightning/__init__.py

@@ -1,18 +1,32 @@
 """CAREamics PyTorch Lightning modules."""
 
 __all__ = [
+    "DataStatsCallback",
     "FCNModule",
     "HyperParametersCallback",
+    "MicroSplitDataModule",
     "PredictDataModule",
     "ProgressBarCallback",
     "TrainDataModule",
     "VAEModule",
     "create_careamics_module",
+    "create_microsplit_predict_datamodule",
+    "create_microsplit_train_datamodule",
     "create_predict_datamodule",
     "create_train_datamodule",
+    "create_unet_based_module",
+    "create_vae_based_module",
 ]
 
-from .callbacks import HyperParametersCallback, ProgressBarCallback
+from .callbacks import DataStatsCallback, HyperParametersCallback, ProgressBarCallback
 from .lightning_module import FCNModule, VAEModule, create_careamics_module
+from .microsplit_data_module import (
+    MicroSplitDataModule,
+    create_microsplit_predict_datamodule,
+    create_microsplit_train_datamodule,
+)
 from .predict_data_module import PredictDataModule, create_predict_datamodule
-from .train_data_module import TrainDataModule, create_train_datamodule
+from .train_data_module import (
+    TrainDataModule,
+    create_train_datamodule,
+)

careamics/lightning/callbacks/__init__.py

@@ -1,11 +1,13 @@
 """Callbacks module."""
 
 __all__ = [
+    "DataStatsCallback",
     "HyperParametersCallback",
     "PredictionWriterCallback",
     "ProgressBarCallback",
 ]
 
+from .data_stats_callback import DataStatsCallback
 from .hyperparameters_callback import HyperParametersCallback
 from .prediction_writer_callback import PredictionWriterCallback
 from .progress_bar_callback import ProgressBarCallback

careamics/lightning/callbacks/data_stats_callback.py

@@ -0,0 +1,23 @@
+"""Data statistics callback."""
+
+import pytorch_lightning as L
+from pytorch_lightning.callbacks import Callback
+
+
+class DataStatsCallback(Callback):
+    """Callback to update model's data statistics from datamodule.
+
+    This callback ensures that the model has access to the data statistics (mean and std)
+    calculated by the datamodule before training starts.
+    """
+
+    def setup(self, trainer: L.Trainer, module: L.LightningModule, stage: str) -> None:
+        """Called when trainer is setting up."""
+        if stage == "fit":
+            # Get data statistics from datamodule
+            (data_mean, data_std), _ = trainer.datamodule.get_data_stats()
+
+            # Set data statistics in the model's likelihood module
+            module.noise_model_likelihood.set_data_stats(
+                data_mean=data_mean["target"], data_std=data_std["target"]
+            )