(Feature) HDN minimum example

src/careamics/careamist.py

@@ -9,11 +9,17 @@
 from pytorch_lightning.callbacks import (
     Callback,
     EarlyStopping,
+    LearningRateMonitor,
     ModelCheckpoint,
 )
 from pytorch_lightning.loggers import CSVLogger, TensorBoardLogger, WandbLogger
 
-from careamics.config import Configuration, UNetBasedAlgorithm, load_configuration
+from careamics.config import (
+    Configuration,
+    UNetBasedAlgorithm,
+    VAEBasedAlgorithm,
+    load_configuration,
+)
 from careamics.config.support import (
     SupportedAlgorithm,
     SupportedArchitecture,
@@ -28,6 +34,7 @@
     PredictDataModule,
     ProgressBarCallback,
     TrainDataModule,
+    VAEModule,
     create_predict_datamodule,
 )
 from careamics.model_io import export_to_bmz, load_pretrained
@@ -141,6 +148,10 @@ def __init__(
                 self.model = FCNModule(
                     algorithm_config=self.cfg.algorithm_config,
                 )
+            elif isinstance(self.cfg.algorithm_config, VAEBasedAlgorithm):
+                self.model = VAEModule(
+                    algorithm_config=self.cfg.algorithm_config,
+                )
             else:
                 raise NotImplementedError("Architecture not supported.")
 
@@ -252,6 +263,7 @@ def _define_callbacks(self, callbacks: Optional[list[Callback]] = None) -> None:
                     **self.cfg.training_config.checkpoint_callback.model_dump(),
                 ),
                 ProgressBarCallback(),
+                LearningRateMonitor(),
             ]
         )
 
@@ -563,6 +575,7 @@ def predict(
         batch_size: int = 1,
         tile_size: Optional[tuple[int, ...]] = None,
         tile_overlap: Optional[tuple[int, ...]] = (48, 48),
+        mmse_count: int = 1,
         axes: Optional[str] = None,
         data_type: Optional[Literal["array", "tiff", "custom"]] = None,
         tta_transforms: bool = False,
@@ -677,6 +690,7 @@ def predict(
         )
 
         # predict
+        self.model.mmse_count = mmse_count  # TODO better way ?
         predictions = self.trainer.predict(
             model=self.model, datamodule=self.pred_datamodule
         )

src/careamics/config/__init__.py

@@ -12,6 +12,7 @@
     "Configuration",
     "DataConfig",
     "GaussianMixtureNMConfig",
+    "HDNAlgorithm",
     "InferenceConfig",
     "LVAELossConfig",
     "MultiChannelNMConfig",
@@ -22,6 +23,7 @@
     "VAEBasedAlgorithm",
     "algorithm_factory",
     "create_care_configuration",
+    "create_hdn_configuration",
     "create_n2n_configuration",
     "create_n2v_configuration",
     "load_configuration",
@@ -30,6 +32,7 @@
 
 from .algorithms import (
     CAREAlgorithm,
+    HDNAlgorithm,
     N2NAlgorithm,
     N2VAlgorithm,
     UNetBasedAlgorithm,
@@ -40,6 +43,7 @@
 from .configuration_factories import (
     algorithm_factory,
     create_care_configuration,
+    create_hdn_configuration,
     create_n2n_configuration,
     create_n2v_configuration,
 )

src/careamics/config/algorithms/__init__.py

@@ -2,13 +2,15 @@
 
 __all__ = [
     "CAREAlgorithm",
+    "HDNAlgorithm",
     "N2NAlgorithm",
     "N2VAlgorithm",
     "UNetBasedAlgorithm",
     "VAEBasedAlgorithm",
 ]
 
 from .care_algorithm_model import CAREAlgorithm
+from .hdn_algorithm_model import HDNAlgorithm
 from .n2n_algorithm_model import N2NAlgorithm
 from .n2v_algorithm_model import N2VAlgorithm
 from .unet_algorithm_model import UNetBasedAlgorithm

src/careamics/config/algorithms/hdn_algorithm_model.py

@@ -0,0 +1,97 @@
+"""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 = "HDN"
+
+HDN_DESCRIPTION = ""
+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
+
+    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

src/careamics/config/algorithms/vae_algorithm_model.py

@@ -40,7 +40,7 @@ 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", "musplit", "denoisplit"]
 
     # NOTE: these are all configs (pydantic models)
     loss: LVAELossConfig
@@ -64,6 +64,14 @@ def algorithm_cross_validation(self: Self) -> Self:
         Self
             The validated model.
         """
+        # hdn
+        if self.algorithm == SupportedAlgorithm.HDN:
+            if self.loss.loss_type != SupportedLoss.HDN:
+                raise ValueError(
+                    f"Algorithm {self.algorithm} only supports loss `hdn`."
+                )
+            if self.model.multiscale_count > 1:
+                raise ValueError("Algorithm `hdn` does not support multiscale models.")
         # musplit
         if self.algorithm == SupportedAlgorithm.MUSPLIT:
             if self.loss.loss_type != SupportedLoss.MUSPLIT:
@@ -108,6 +116,12 @@ def output_channels_validation(self: Self) -> Self:
                 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,15 @@ def predict_logvar_validation(self: Self) -> Self:
                 "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`."
+        #         )
         return self
 
     def __str__(self) -> str:

src/careamics/config/architectures/lvae_model.py

@@ -15,11 +15,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: list[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 ?
@@ -126,6 +124,13 @@ def validate_input_shape(cls, input_shape: list) -> list:
                 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 +260,4 @@ def is_3D(self) -> bool:
         bool
             Whether the model is 3D or not.
         """
-        return self.conv_dims == 3
+        return len(self.input_shape) == 3

src/careamics/config/configuration.py

@@ -12,6 +12,7 @@
 
 from careamics.config.algorithms import (
     CAREAlgorithm,
+    HDNAlgorithm,
     N2NAlgorithm,
     N2VAlgorithm,
 )
@@ -22,6 +23,7 @@
     CAREAlgorithm,
     N2NAlgorithm,
     N2VAlgorithm,
+    HDNAlgorithm,
 ]