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

careamics/config/algorithms/microsplit_algorithm_model.py

@@ -0,0 +1,103 @@
+"""MicroSplit 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
+
+MICROSPLIT = "MicroSplit"
+
+MICROSPLIT_DESCRIPTION = """MicroSplit is a self-supervised deep learning method for
+microscopy image splitting that combines the strengths of both denoising and
+representation learning approaches."""
+
+MICROSPLIT_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 MicroSplitAlgorithm(VAEBasedAlgorithm):
+    """MicroSplit algorithm configuration."""
+
+    model_config = ConfigDict(validate_assignment=True)
+
+    algorithm: Literal["microsplit"] = "microsplit"
+
+    loss: LVAELossConfig
+
+    model: LVAEModel  # TODO add validators
+
+    is_supervised: bool = True
+
+    def get_algorithm_friendly_name(self) -> str:
+        """
+        Get the algorithm friendly name.
+
+        Returns
+        -------
+        str
+            Friendly name of the algorithm.
+        """
+        return MICROSPLIT
+
+    def get_algorithm_keywords(self) -> list[str]:
+        """
+        Get algorithm keywords.
+
+        Returns
+        -------
+        list[str]
+            List of keywords.
+        """
+        return [
+            "restoration",
+            "VAE",
+            "self-supervised",
+            "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 MICROSPLIT_REF.text + " doi: " + MICROSPLIT_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 [MICROSPLIT_REF]
+
+    def get_algorithm_description(self) -> str:
+        """
+        Get the algorithm description.
+
+        Returns
+        -------
+        str
+            Algorithm description.
+        """
+        return MICROSPLIT_DESCRIPTION

careamics/config/algorithms/n2v_algorithm_model.py

@@ -1,10 +1,9 @@
 """N2V Algorithm configuration."""
 
-from typing import Annotated, Literal
+from typing import Annotated, Literal, Self
 
 from bioimageio.spec.generic.v0_3 import CiteEntry
 from pydantic import AfterValidator, ConfigDict, model_validator
-from typing_extensions import Self
 
 from careamics.config.architectures import UNetModel
 from careamics.config.support import SupportedPixelManipulation, SupportedStructAxis

careamics/config/algorithms/vae_algorithm_model.py

@@ -3,10 +3,9 @@
 from __future__ import annotations
 
 from pprint import pformat
-from typing import Literal, Optional
+from typing import Literal, Self
 
 from pydantic import BaseModel, ConfigDict, model_validator
-from typing_extensions import Self
 
 from careamics.config.architectures import LVAEModel
 from careamics.config.likelihood_model import (
@@ -40,14 +39,17 @@ class VAEBasedAlgorithm(BaseModel):
     # defined in SupportedAlgorithm
     # TODO: Use supported Enum classes for typing?
     #   - values can still be passed as strings and they will be cast to Enum
-    algorithm: Literal["musplit", "denoisplit"]
+    algorithm: Literal["hdn", "microsplit"]
 
     # NOTE: these are all configs (pydantic models)
     loss: LVAELossConfig
     model: LVAEModel
-    noise_model: Optional[MultiChannelNMConfig] = None
-    noise_model_likelihood: Optional[NMLikelihoodConfig] = None
-    gaussian_likelihood: Optional[GaussianLikelihoodConfig] = None
+    noise_model: MultiChannelNMConfig | None = None
+    noise_model_likelihood: NMLikelihoodConfig | None = None
+    gaussian_likelihood: GaussianLikelihoodConfig | None = None  # TODO change to str
+
+    mmse_count: int = 1
+    is_supervised: bool = False
 
     # Optional fields
     optimizer: OptimizerModel = OptimizerModel()
@@ -64,22 +66,26 @@ class VAEBasedAlgorithm(BaseModel):
         Self
             The validated model.
         """
-        # musplit
-        if self.algorithm == SupportedAlgorithm.MUSPLIT:
-            if self.loss.loss_type != SupportedLoss.MUSPLIT:
+        # hdn
+        # TODO move to designated configurations
+        if self.algorithm == SupportedAlgorithm.HDN:
+            if self.loss.loss_type != SupportedLoss.HDN:
                 raise ValueError(
-                    f"Algorithm {self.algorithm} only supports loss `musplit`."
+                    f"Algorithm {self.algorithm} only supports loss `hdn`."
                 )
-
-        if self.algorithm == SupportedAlgorithm.DENOISPLIT:
+            if self.model.multiscale_count > 1:
+                raise ValueError("Algorithm `hdn` does not support multiscale models.")
+        # musplit
+        if self.algorithm == SupportedAlgorithm.MICROSPLIT:
             if self.loss.loss_type not in [
+                SupportedLoss.MUSPLIT,
                 SupportedLoss.DENOISPLIT,
                 SupportedLoss.DENOISPLIT_MUSPLIT,
-            ]:
+            ]:  # TODO Update losses configs, make loss just microsplit
                 raise ValueError(
-                    f"Algorithm {self.algorithm} only supports loss `denoisplit` "
-                    "or `denoisplit_musplit."
-                )
+                    f"Algorithm {self.algorithm} only supports loss `microsplit`."
+                )  # TODO Update losses configs
+
             if (
                 self.loss.loss_type == SupportedLoss.DENOISPLIT
                 and self.model.predict_logvar is not None
@@ -88,8 +94,10 @@ class VAEBasedAlgorithm(BaseModel):
                     "Algorithm `denoisplit` with loss `denoisplit` only supports "
                     "`predict_logvar` as `None`."
                 )
-
-            if self.noise_model is None:
+            if (
+                self.loss.loss_type == SupportedLoss.DENOISPLIT
+                and self.noise_model is None
+            ):
                 raise ValueError("Algorithm `denoisplit` requires a noise model.")
         # TODO: what if algorithm is not musplit or denoisplit
         return self
@@ -108,6 +116,12 @@ class VAEBasedAlgorithm(BaseModel):
                 f"Number of output channels ({self.model.output_channels}) must match "
                 f"the number of noise models ({len(self.noise_model.noise_models)})."
             )
+
+        if self.algorithm == SupportedAlgorithm.HDN:
+            assert self.model.output_channels == 1, (
+                f"Number of output channels ({self.model.output_channels}) must be 1 "
+                "for algorithm `hdn`."
+            )
         return self
 
     @model_validator(mode="after")
@@ -127,6 +141,16 @@ class VAEBasedAlgorithm(BaseModel):
                 "Gaussian likelihood model `predict_logvar` "
                 f"({self.gaussian_likelihood.predict_logvar}).",
             )
+        # if self.algorithm == SupportedAlgorithm.HDN:
+        #     assert (
+        #         self.model.predict_logvar is None
+        #     ), "Model `predict_logvar` must be `None` for algorithm `hdn`."
+        #     if self.gaussian_likelihood is not None:
+        #         assert self.gaussian_likelihood.predict_logvar is None, (
+        #             "Gaussian likelihood model `predict_logvar` must be `None` "
+        #             "for algorithm `hdn`."
+        #         )
+        # TODO check this
         return self
 
     def __str__(self) -> str:
@@ -138,3 +162,14 @@ class VAEBasedAlgorithm(BaseModel):
             Pretty string.
         """
         return pformat(self.model_dump())
+
+    @classmethod
+    def get_compatible_algorithms(cls) -> list[str]:
+        """Get the list of compatible algorithms.
+
+        Returns
+        -------
+        list of str
+            List of compatible algorithms.
+        """
+        return ["hdn", "microsplit"]

careamics/config/architectures/lvae_model.py

@@ -1,9 +1,8 @@
 """LVAE Pydantic model."""
 
-from typing import Literal
+from typing import Literal, Self
 
 from pydantic import ConfigDict, Field, field_validator, model_validator
-from typing_extensions import Self
 
 from .architecture_model import ArchitectureModel
 
@@ -15,11 +14,9 @@ class LVAEModel(ArchitectureModel):
     model_config = ConfigDict(validate_assignment=True, validate_default=True)
 
     architecture: Literal["LVAE"]
-    """Name of the architecture."""
-
-    input_shape: list[int] = Field(default=[64, 64], validate_default=True)
-    """Shape of the input patch (C, Z, Y, X) or (C, Y, X) if the data is 2D."""
 
+    input_shape: tuple[int, ...] = Field(default=(64, 64), validate_default=True)
+    """Shape of the input patch (Z, Y, X) or (Y, X) if the data is 2D."""
     encoder_conv_strides: list = Field(default=[2, 2], validate_default=True)
 
     # TODO make this per hierarchy step ?
@@ -42,7 +39,7 @@ class LVAEModel(ArchitectureModel):
         default="ELU",
     )
 
-    predict_logvar: Literal[None, "pixelwise"] = None
+    predict_logvar: Literal[None, "pixelwise"] = "pixelwise"
     analytical_kl: bool = Field(default=False)
 
     @model_validator(mode="after")
@@ -126,6 +123,13 @@ class LVAEModel(ArchitectureModel):
                 f"Input shape must be greater than 1 in all dimensions"
                 f"(got {input_shape})."
             )
+
+        if any(s < 64 for s in input_shape[-2:]):
+            raise ValueError(
+                f"Input shape must be greater or equal to 64 in XY dimensions"
+                f"(got {input_shape})."
+            )
+
         return input_shape
 
     @field_validator("encoder_n_filters")
@@ -255,4 +259,4 @@ class LVAEModel(ArchitectureModel):
         bool
             Whether the model is 3D or not.
         """
-        return self.conv_dims == 3
+        return len(self.input_shape) == 3

careamics/config/callback_model.py

@@ -3,7 +3,7 @@
 from __future__ import annotations
 
 from datetime import timedelta
-from typing import Literal, Optional
+from typing import Literal
 
 from pydantic import (
     BaseModel,
@@ -24,7 +24,7 @@ class CheckpointModel(BaseModel):
 
     model_config = ConfigDict(validate_assignment=True, validate_default=True)
 
-    monitor: Literal["val_loss"] = Field(default="val_loss")
+    monitor: Literal["val_loss"] | str | None = Field(default="val_loss")
     """Quantity to monitor, currently only `val_loss`."""
 
     verbose: bool = Field(default=False)
@@ -33,11 +33,15 @@ class CheckpointModel(BaseModel):
     save_weights_only: bool = Field(default=False)
     """When `True`, only the model's weights will be saved (model.save_weights)."""
 
-    save_last: Optional[Literal[True, False, "link"]] = Field(default=True)
+    save_last: Literal[True, False, "link"] | None = Field(default=True)
     """When `True`, saves a last.ckpt copy whenever a checkpoint file gets saved."""
 
-    save_top_k: int = Field(default=3, ge=-1, le=100)
-    """If `save_top_k == kz, the best k models according to the quantity monitored
+    save_top_k: int = Field(
+        default=3,
+        ge=-1,
+        le=100,
+    )
+    """If `save_top_k == k, the best k models according to the quantity monitored
     will be saved. If `save_top_k == 0`, no models are saved. if `save_top_k == -1`,
     all models are saved."""
 
@@ -51,13 +55,13 @@ class CheckpointModel(BaseModel):
     auto_insert_metric_name: bool = Field(default=False)
     """When `True`, the checkpoints filenames will contain the metric name."""
 
-    every_n_train_steps: Optional[int] = Field(default=None, ge=1, le=1000)
+    every_n_train_steps: int | None = Field(default=None, ge=1, le=1000)
     """Number of training steps between checkpoints."""
 
-    train_time_interval: Optional[timedelta] = Field(default=None)
+    train_time_interval: timedelta | None = Field(default=None)
     """Checkpoints are monitored at the specified time interval."""
 
-    every_n_epochs: Optional[int] = Field(default=None, ge=1, le=100)
+    every_n_epochs: int | None = Field(default=None, ge=1, le=100)
     """Number of epochs between checkpoints."""
 
 
@@ -96,14 +100,14 @@ class EarlyStoppingModel(BaseModel):
     """When `True`, stops training when the monitored quantity becomes `NaN` or
     `inf`."""
 
-    stopping_threshold: Optional[float] = Field(default=None)
+    stopping_threshold: float | None = Field(default=None)
     """Stop training immediately once the monitored quantity reaches this threshold."""
 
-    divergence_threshold: Optional[float] = Field(default=None)
+    divergence_threshold: float | None = Field(default=None)
     """Stop training as soon as the monitored quantity becomes worse than this
     threshold."""
 
-    check_on_train_epoch_end: Optional[bool] = Field(default=False)
+    check_on_train_epoch_end: bool | None = Field(default=False)
     """Whether to run early stopping at the end of the training epoch. If this is
     `False`, then the check runs at the end of the validation."""
 

careamics/config/configuration.py

@@ -5,24 +5,28 @@ from __future__ import annotations
 import re
 from collections.abc import Callable
 from pprint import pformat
-from typing import Any, Literal, Union
+from typing import Any, Literal, Self, Union
 
 import numpy as np
 from bioimageio.spec.generic.v0_3 import CiteEntry
 from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator
 from pydantic.main import IncEx
-from typing_extensions import Self
 
 from careamics.config.algorithms import (
     CAREAlgorithm,
+    HDNAlgorithm,
+    MicroSplitAlgorithm,
     N2NAlgorithm,
     N2VAlgorithm,
 )
 from careamics.config.data import DataConfig
 from careamics.config.training_model import TrainingConfig
+from careamics.lvae_training.dataset.config import MicroSplitDataConfig
 
 ALGORITHMS = Union[
     CAREAlgorithm,
+    HDNAlgorithm,
+    MicroSplitAlgorithm,
     N2NAlgorithm,
     N2VAlgorithm,
 ]
@@ -86,7 +90,6 @@ class Configuration(BaseModel):
     ...     axes="YX",
     ...     patch_size=[64, 64],
     ...     batch_size=32,
-    ...     num_epochs=100
     ... )
 
     The configuration can be exported to a dictionary using the model_dump method:
@@ -110,9 +113,7 @@ class Configuration(BaseModel):
     ...                 "architecture": "UNet",
     ...             },
     ...         },
-    ...         "training_config": {
-    ...             "num_epochs": 200,
-    ...         },
+    ...         "training_config": {},
     ...         "data_config": {
     ...             "data_type": "tiff",
     ...             "patch_size": [64, 64],
@@ -140,7 +141,7 @@ class Configuration(BaseModel):
     """Algorithm configuration, holding all parameters required to configure the
     model."""
 
-    data_config: DataConfig
+    data_config: DataConfig | MicroSplitDataConfig
     """Data configuration, holding all parameters required to configure the training
     data loader."""
 
@@ -185,7 +186,7 @@ class Configuration(BaseModel):
 
         return name
 
-    @model_validator(mode="after")
+    @model_validator(mode="after")  # TODO move to n2v configs or remove
     def validate_n2v_mask_pixel_perc(self: Self) -> Self:
         """
         Validate that there will always be at least one blind-spot pixel in every patch.