careamics 0.0.15__py3-none-any.whl → 0.0.17__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
@@ -208,13 +207,12 @@ class DataConfig(BaseModel):
 
     @field_validator("train_dataloader_params", "val_dataloader_params", mode="before")
     @classmethod
-    def set_default_dataloader_params(
+    def set_default_pin_memory(
         cls, dataloader_params: dict[str, Any]
     ) -> dict[str, Any]:
         """
-        Set default dataloader parameters if not provided.
+        Set default pin_memory for dataloader parameters if not provided.
 
-        - If 'num_workers' is not set, it defaults to the number of available CPU cores.
         - If 'pin_memory' is not set, it defaults to True if CUDA is available.
 
         Parameters
@@ -225,21 +223,62 @@ class DataConfig(BaseModel):
         Returns
         -------
         dict of {str: Any}
-            The dataloader parameters with defaults applied.
+            The dataloader parameters with pin_memory default applied.
+        """
+        if "pin_memory" not in dataloader_params:
+            import torch
+
+            dataloader_params["pin_memory"] = torch.cuda.is_available()
+
+        return dataloader_params
+
+    @field_validator("train_dataloader_params", mode="before")
+    @classmethod
+    def set_default_train_workers(
+        cls, dataloader_params: dict[str, Any]
+    ) -> dict[str, Any]:
+        """
+        Set default num_workers for training dataloader if not provided.
+
+        - If 'num_workers' is not set, it defaults to the number of available CPU cores.
+
+        Parameters
+        ----------
+        dataloader_params : dict of {str: Any}
+            The training dataloader parameters.
+
+        Returns
+        -------
+        dict of {str: Any}
+            The dataloader parameters with num_workers default applied.
         """
         if "num_workers" not in dataloader_params:
-            # Use 1 worker during tests, otherwise use all available CPU cores
+            # Use 0 workers during tests, otherwise use all available CPU cores
             if "pytest" in sys.modules:
                 dataloader_params["num_workers"] = 0
             else:
                 dataloader_params["num_workers"] = os.cpu_count()
 
-        if "pin_memory" not in dataloader_params:
-            import torch
+        return dataloader_params
 
-            dataloader_params["pin_memory"] = torch.cuda.is_available()
+    @model_validator(mode="after")
+    def set_val_workers_to_match_train(self: Self) -> Self:
+        """
+        Set validation dataloader num_workers to match training dataloader.
 
-        return dataloader_params
+        If num_workers is not specified in val_dataloader_params, it will be set to the
+        same value as train_dataloader_params["num_workers"].
+
+        Returns
+        -------
+        Self
+            Validated data model with synchronized num_workers.
+        """
+        if "num_workers" not in self.val_dataloader_params:
+            self.val_dataloader_params["num_workers"] = self.train_dataloader_params[
+                "num_workers"
+            ]
+        return self
 
     @field_validator("train_dataloader_params")
     @classmethod

careamics/config/data/ng_data_model.py

@@ -2,9 +2,12 @@
 
 from __future__ import annotations
 
+import os
+import random
+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
@@ -17,10 +20,15 @@ from pydantic import (
     field_validator,
     model_validator,
 )
-from typing_extensions import Self
 
 from ..transformations import XYFlipModel, XYRandomRotate90Model
 from ..validators import check_axes_validity
+from .patch_filter import (
+    MaskFilterModel,
+    MaxFilterModel,
+    MeanSTDFilterModel,
+    ShannonFilterModel,
+)
 from .patching_strategies import (
     RandomPatchingModel,
     TiledPatchingModel,
@@ -38,6 +46,17 @@ from .patching_strategies import (
 #       - or is the responsibility of the creator (e.g. conveneince functions)
 
 
+def generate_random_seed() -> int:
+    """Generate a random seed for reproducibility.
+
+    Returns
+    -------
+    int
+        A random integer between 1 and 2^31 - 1.
+    """
+    return random.randint(1, 2**31 - 1)
+
+
 def np_float_to_scientific_str(x: float) -> str:
     """Return a string scientific representation of a float.
 
@@ -68,6 +87,16 @@ PatchingStrategies = Union[
 ]
 """Patching strategies."""
 
+PatchFilters = Union[
+    MaxFilterModel,
+    MeanSTDFilterModel,
+    ShannonFilterModel,
+]
+"""Patch filters."""
+
+CoordFilters = Union[MaskFilterModel]  # add more here as needed
+"""Coordinate filters."""
+
 
 class NGDataConfig(BaseModel):
     """Next-Generation Dataset configuration.
@@ -106,6 +135,18 @@ class NGDataConfig(BaseModel):
     batch_size: int = Field(default=1, ge=1, validate_default=True)
     """Batch size for training."""
 
+    patch_filter: PatchFilters | None = Field(default=None, discriminator="name")
+    """Patch filter to apply when using random patching. Only available during
+    training."""
+
+    coord_filter: CoordFilters | None = Field(default=None, discriminator="name")
+    """Coordinate filter to apply when using random patching. Only available during
+    training."""
+
+    patch_filter_patience: int = Field(default=5, ge=1)
+    """Number of consecutive patches not passing the filter before accepting the next
+    patch."""
+
     image_means: list[Float] | None = Field(default=None, min_length=0, max_length=32)
     """Means of the data across channels, used for normalization."""
 
@@ -142,8 +183,8 @@ class NGDataConfig(BaseModel):
     test_dataloader_params: dict[str, Any] = Field(default={})
     """Dictionary of PyTorch test dataloader parameters."""
 
-    seed: int | None = Field(default=None, gt=0)
-    """Random seed for reproducibility."""
+    seed: int | None = Field(default_factory=generate_random_seed, gt=0)
+    """Random seed for reproducibility. If not specified, a random seed is generated."""
 
     @field_validator("axes")
     @classmethod
@@ -297,6 +338,129 @@ class NGDataConfig(BaseModel):
 
         return self
 
+    @model_validator(mode="after")
+    def propagate_seed_to_filters(self: Self) -> Self:
+        """
+        Propagate the main seed to patch and coordinate filters that support seeds.
+
+        This ensures that all filters use the same seed for reproducibility,
+        unless they already have a seed explicitly set.
+
+        Returns
+        -------
+        Self
+            Data model with propagated seeds.
+        """
+        if self.seed is not None:
+            if self.patch_filter is not None:
+                if (
+                    hasattr(self.patch_filter, "seed")
+                    and self.patch_filter.seed is None
+                ):
+                    self.patch_filter.seed = self.seed
+
+            if self.coord_filter is not None:
+                if (
+                    hasattr(self.coord_filter, "seed")
+                    and self.coord_filter.seed is None
+                ):
+                    self.coord_filter.seed = self.seed
+
+        return self
+
+    @model_validator(mode="after")
+    def propagate_seed_to_transforms(self: Self) -> Self:
+        """
+        Propagate the main seed to all transforms that support seeds.
+
+        This ensures that all transforms use the same seed for reproducibility,
+        unless they already have a seed explicitly set.
+
+        Returns
+        -------
+        Self
+            Data model with propagated seeds.
+        """
+        if self.seed is not None:
+            for transform in self.transforms:
+                if hasattr(transform, "seed") and transform.seed is None:
+                    transform.seed = self.seed
+        return self
+
+    @field_validator("train_dataloader_params", "val_dataloader_params", mode="before")
+    @classmethod
+    def set_default_pin_memory(
+        cls, dataloader_params: dict[str, Any]
+    ) -> dict[str, Any]:
+        """
+        Set default pin_memory for dataloader parameters if not provided.
+
+        - If 'pin_memory' is not set, it defaults to True if CUDA is available.
+
+        Parameters
+        ----------
+        dataloader_params : dict of {str: Any}
+            The dataloader parameters.
+
+        Returns
+        -------
+        dict of {str: Any}
+            The dataloader parameters with pin_memory default applied.
+        """
+        if "pin_memory" not in dataloader_params:
+            import torch
+
+            dataloader_params["pin_memory"] = torch.cuda.is_available()
+        return dataloader_params
+
+    @field_validator("train_dataloader_params", mode="before")
+    @classmethod
+    def set_default_train_workers(
+        cls, dataloader_params: dict[str, Any]
+    ) -> dict[str, Any]:
+        """
+        Set default num_workers for training dataloader if not provided.
+
+        - If 'num_workers' is not set, it defaults to the number of available CPU cores.
+
+        Parameters
+        ----------
+        dataloader_params : dict of {str: Any}
+            The training dataloader parameters.
+
+        Returns
+        -------
+        dict of {str: Any}
+            The dataloader parameters with num_workers default applied.
+        """
+        if "num_workers" not in dataloader_params:
+            # Use 0 workers during tests, otherwise use all available CPU cores
+            if "pytest" in sys.modules:
+                dataloader_params["num_workers"] = 0
+            else:
+                dataloader_params["num_workers"] = os.cpu_count()
+
+        return dataloader_params
+
+    @model_validator(mode="after")
+    def set_val_workers_to_match_train(self: Self) -> Self:
+        """
+        Set validation dataloader num_workers to match training dataloader.
+
+        If num_workers is not specified in val_dataloader_params, it will be set to the
+        same value as train_dataloader_params["num_workers"].
+
+        Returns
+        -------
+        Self
+            Validated data model with synchronized num_workers.
+        """
+        if "num_workers" not in self.val_dataloader_params:
+            self.val_dataloader_params["num_workers"] = self.train_dataloader_params[
+                "num_workers"
+            ]
+        return self
+
     def __str__(self) -> str:
         """
         Pretty string reprensenting the configuration.

careamics/config/data/patch_filter/__init__.py

@@ -0,0 +1,15 @@
+"""Pydantic models representing coordinate and patch filters."""
+
+__all__ = [
+    "FilterModel",
+    "MaskFilterModel",
+    "MaxFilterModel",
+    "MeanSTDFilterModel",
+    "ShannonFilterModel",
+]
+
+from .filter_model import FilterModel
+from .mask_filter_model import MaskFilterModel
+from .max_filter_model import MaxFilterModel
+from .meanstd_filter_model import MeanSTDFilterModel
+from .shannon_filter_model import ShannonFilterModel

careamics/config/data/patch_filter/filter_model.py

@@ -0,0 +1,16 @@
+"""Base class for patch and coordinate filtering models."""
+
+from pydantic import BaseModel, Field
+
+
+class FilterModel(BaseModel):
+    """Base class for patch and coordinate filtering models."""
+
+    name: str
+    """Name of the filter."""
+
+    p: float = Field(1.0, ge=0.0, le=1.0)
+    """Probability of applying the filter to a patch or coordinate."""
+
+    seed: int | None = Field(default=None, gt=0)
+    """Seed for the random number generator for reproducibility."""

careamics/config/data/patch_filter/mask_filter_model.py

@@ -0,0 +1,17 @@
+"""Pydantic model for the mask coordinate filter."""
+
+from typing import Literal
+
+from pydantic import Field
+
+from .filter_model import FilterModel
+
+
+class MaskFilterModel(FilterModel):
+    """Pydantic model for the mask coordinate filter."""
+
+    name: Literal["mask"] = "mask"
+    """Name of the filter."""
+
+    coverage: float = Field(0.5, ge=0.0, le=1.0)
+    """Percentage of masked pixels required to keep a patch."""

careamics/config/data/patch_filter/max_filter_model.py

@@ -0,0 +1,15 @@
+"""Pydantic model for the max patch filter."""
+
+from typing import Literal
+
+from .filter_model import FilterModel
+
+
+class MaxFilterModel(FilterModel):
+    """Pydantic model for the max patch filter."""
+
+    name: Literal["max"] = "max"
+    """Name of the filter."""
+
+    threshold: float
+    """Threshold for the minimum of the max-filtered patch."""

careamics/config/data/patch_filter/meanstd_filter_model.py

@@ -0,0 +1,18 @@
+"""Pydantic model for the mean std patch filter."""
+
+from typing import Literal
+
+from .filter_model import FilterModel
+
+
+class MeanSTDFilterModel(FilterModel):
+    """Pydantic model for the mean std patch filter."""
+
+    name: Literal["mean_std"] = "mean_std"
+    """Name of the filter."""
+
+    mean_threshold: float
+    """Minimum mean intensity required to keep a patch."""
+
+    std_threshold: float | None = None
+    """Minimum standard deviation required to keep a patch."""

careamics/config/data/patch_filter/shannon_filter_model.py

@@ -0,0 +1,15 @@
+"""Pydantic model for the Shannon entropy patch filter."""
+
+from typing import Literal
+
+from .filter_model import FilterModel
+
+
+class ShannonFilterModel(FilterModel):
+    """Pydantic model for the Shannon entropy patch filter."""
+
+    name: Literal["shannon"] = "shannon"
+    """Name of the filter."""
+
+    threshold: float
+    """Minimum Shannon entropy required to keep a patch."""

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_filters.py

@@ -0,0 +1,17 @@
+"""Coordinate and patch filters supported by CAREamics."""
+
+from careamics.utils import BaseEnum
+
+
+class SupportedPatchFilters(str, BaseEnum):
+    """Supported patch filters."""
+
+    MAX = "max"
+    MEANSTD = "mean_std"
+    SHANNON = "shannon"
+
+
+class SupportedCoordinateFilters(str, BaseEnum):
+    """Supported coordinate filters."""
+
+    MASK = "mask"

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