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

careamics/lightning/train_data_module.py

@@ -2,7 +2,7 @@
 
 from collections.abc import Callable
 from pathlib import Path
-from typing import Any, Literal, Optional, Union
+from typing import Any, Literal, Union
 
 import numpy as np
 import pytorch_lightning as L
@@ -121,10 +121,10 @@ class TrainDataModule(L.LightningDataModule):
         self,
         data_config: DataConfig,
         train_data: Union[Path, str, NDArray],
-        val_data: Optional[Union[Path, str, NDArray]] = None,
-        train_data_target: Optional[Union[Path, str, NDArray]] = None,
-        val_data_target: Optional[Union[Path, str, NDArray]] = None,
-        read_source_func: Optional[Callable] = None,
+        val_data: Union[Path, str, NDArray] | None = None,
+        train_data_target: Union[Path, str, NDArray] | None = None,
+        val_data_target: Union[Path, str, NDArray] | None = None,
+        read_source_func: Callable | None = None,
         extension_filter: str = "",
         val_percentage: float = 0.1,
         val_minimum_split: int = 5,
@@ -477,15 +477,16 @@ def create_train_datamodule(
     patch_size: list[int],
     axes: str,
     batch_size: int,
-    val_data: Optional[Union[str, Path, NDArray]] = None,
-    transforms: Optional[list[TransformModel]] = None,
-    train_target_data: Optional[Union[str, Path, NDArray]] = None,
-    val_target_data: Optional[Union[str, Path, NDArray]] = None,
-    read_source_func: Optional[Callable] = None,
+    val_data: Union[str, Path, NDArray] | None = None,
+    transforms: list[TransformModel] | None = None,
+    train_target_data: Union[str, Path, NDArray] | None = None,
+    val_target_data: Union[str, Path, NDArray] | None = None,
+    read_source_func: Callable | None = None,
     extension_filter: str = "",
     val_percentage: float = 0.1,
     val_minimum_patches: int = 5,
-    dataloader_params: Optional[dict] = None,
+    train_dataloader_params: dict | None = None,
+    val_dataloader_params: dict | None = None,
     use_in_memory: bool = True,
 ) -> TrainDataModule:
     """Create a TrainDataModule.
@@ -556,8 +557,10 @@ def create_train_datamodule(
     val_minimum_patches : int, optional
         Minimum number of patches to split from the training data for validation if
         no validation data is given, by default 5.
-    dataloader_params : dict, optional
-        Pytorch dataloader parameters, by default {}.
+    train_dataloader_params : dict, optional
+        Pytorch dataloader parameters for the training data, by default {}.
+    val_dataloader_params : dict, optional
+        Pytorch dataloader parameters for the validation data, by default {}.
     use_in_memory : bool, optional
         Use in memory dataset if possible, by default True.
 
@@ -617,8 +620,11 @@ def create_train_datamodule(
     ...     transforms=my_transforms,
     ... )
     """
-    if dataloader_params is None:
-        dataloader_params = {}
+    if train_dataloader_params is None:
+        train_dataloader_params = {"shuffle": True}
+
+    if val_dataloader_params is None:
+        val_dataloader_params = {"shuffle": False}
 
     data_dict: dict[str, Any] = {
         "mode": "train",
@@ -626,7 +632,8 @@ def create_train_datamodule(
         "patch_size": patch_size,
         "axes": axes,
         "batch_size": batch_size,
-        "dataloader_params": dataloader_params,
+        "train_dataloader_params": train_dataloader_params,
+        "val_dataloader_params": val_dataloader_params,
     }
 
     # if transforms are passed (otherwise it will use the default ones)
@@ -637,9 +644,13 @@ def create_train_datamodule(
     data_config = DataConfig(**data_dict)
 
     # sanity check on the dataloader parameters
-    if "batch_size" in dataloader_params:
+    if "batch_size" in train_dataloader_params:
+        # remove it
+        del train_dataloader_params["batch_size"]
+
+    if "batch_size" in val_dataloader_params:
         # remove it
-        del dataloader_params["batch_size"]
+        del val_dataloader_params["batch_size"]
 
     return TrainDataModule(
         data_config=data_config,

careamics/losses/__init__.py

@@ -3,6 +3,7 @@
 __all__ = [
     "denoisplit_loss",
     "denoisplit_musplit_loss",
+    "hdn_loss",
     "loss_factory",
     "mae_loss",
     "mse_loss",
@@ -12,4 +13,9 @@ __all__ = [
 
 from .fcn.losses import mae_loss, mse_loss, n2v_loss
 from .loss_factory import loss_factory
-from .lvae.losses import denoisplit_loss, denoisplit_musplit_loss, musplit_loss
+from .lvae.losses import (
+    denoisplit_loss,
+    denoisplit_musplit_loss,
+    hdn_loss,
+    musplit_loss,
+)

careamics/losses/loss_factory.py

@@ -14,7 +14,12 @@ from torch import Tensor as tensor
 
 from ..config.support import SupportedLoss
 from .fcn.losses import mae_loss, mse_loss, n2v_loss
-from .lvae.losses import denoisplit_loss, denoisplit_musplit_loss, musplit_loss
+from .lvae.losses import (
+    denoisplit_loss,
+    denoisplit_musplit_loss,
+    hdn_loss,
+    musplit_loss,
+)
 
 
 @dataclass
@@ -59,6 +64,9 @@ def loss_factory(loss: Union[SupportedLoss, str]) -> Callable:
     elif loss == SupportedLoss.MSE:
         return mse_loss
 
+    elif loss == SupportedLoss.HDN:
+        return hdn_loss
+
     elif loss == SupportedLoss.MUSPLIT:
         return musplit_loss
 

careamics/losses/lvae/losses.py

@@ -2,7 +2,7 @@
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Any, Literal, Optional, Union
+from typing import TYPE_CHECKING, Any, Literal, Union
 
 import numpy as np
 import torch
@@ -89,6 +89,7 @@ def _reconstruction_loss_musplit_denoisplit(
     if predictions.shape[1] == 2 * targets.shape[1]:
         # predictions contain both mean and log-variance
         pred_mean, _ = predictions.chunk(2, dim=1)
+        # TODO if this condition does not hold, everything breaks later!
     else:
         pred_mean = predictions
 
@@ -112,7 +113,7 @@ def get_kl_divergence_loss(
     rescaling: Literal["latent_dim", "image_dim"],
     aggregation: Literal["mean", "sum"],
     free_bits_coeff: float,
-    img_shape: Optional[tuple[int]] = None,
+    img_shape: tuple[int] | None = None,
 ) -> torch.Tensor:
     """Compute the KL divergence loss.
 
@@ -269,13 +270,97 @@ def _get_kl_divergence_loss_denoisplit(
 # - `__init__` method initializes the loss parameters now contained in
 # the `LVAELossParameters` class
 # NOTE: same for the other loss functions
+
+
+def hdn_loss(
+    model_outputs: tuple[torch.Tensor, dict[str, Any]],
+    targets: torch.Tensor,
+    config: LVAELossConfig,
+    gaussian_likelihood: GaussianLikelihood | None,
+    noise_model_likelihood: NoiseModelLikelihood | None,
+) -> dict[str, torch.Tensor] | None:
+    """Loss function for HDN.
+
+    Parameters
+    ----------
+    model_outputs : tuple[torch.Tensor, dict[str, Any]]
+        Tuple containing the model predictions (shape is (B, `target_ch`, [Z], Y, X))
+        and the top-down layer data (e.g., sampled latents, KL-loss values, etc.).
+    targets : torch.Tensor
+        The target image used to compute the reconstruction loss. In this case we use
+        the input patch itself as target. Shape is (B, `target_ch`, [Z], Y, X).
+    config : LVAELossConfig
+        The config for loss function containing all loss hyperparameters.
+    gaussian_likelihood : GaussianLikelihood
+        The Gaussian likelihood object.
+    noise_model_likelihood : NoiseModelLikelihood
+        The noise model likelihood object.
+
+    Returns
+    -------
+    output : Optional[dict[str, torch.Tensor]]
+        A dictionary containing the overall loss `["loss"]`, the reconstruction loss
+        `["reconstruction_loss"]`, and the KL divergence loss `["kl_loss"]`.
+    """
+    if gaussian_likelihood is not None:
+        likelihood = gaussian_likelihood
+    elif noise_model_likelihood is not None:
+        likelihood = noise_model_likelihood
+    else:
+        raise ValueError("Invalid likelihood object.")
+    # TODO refactor loss signature
+    predictions, td_data = model_outputs
+
+    # Reconstruction loss computation
+    recons_loss = config.reconstruction_weight * get_reconstruction_loss(
+        reconstruction=predictions,
+        target=targets,
+        likelihood_obj=likelihood,
+    )
+    if torch.isnan(recons_loss).any():
+        recons_loss = 0.0
+
+    # KL loss computation
+    kl_weight = get_kl_weight(
+        config.kl_params.annealing,
+        config.kl_params.start,
+        config.kl_params.annealtime,
+        config.kl_weight,
+        config.kl_params.current_epoch,
+    )
+    kl_loss = (
+        _get_kl_divergence_loss_denoisplit(
+            topdown_data=td_data,
+            img_shape=targets.shape[2:],
+            kl_type=config.kl_params.loss_type,
+        )
+        * kl_weight
+    )
+
+    net_loss = recons_loss + kl_loss  # TODO add check that losses coefs sum to 1
+    output = {
+        "loss": net_loss,
+        "reconstruction_loss": (
+            recons_loss.detach()
+            if isinstance(recons_loss, torch.Tensor)
+            else recons_loss
+        ),
+        "kl_loss": kl_loss.detach(),
+    }
+    # https://github.com/openai/vdvae/blob/main/train.py#L26
+    if torch.isnan(net_loss).any():
+        return None
+
+    return output
+
+
 def musplit_loss(
     model_outputs: tuple[torch.Tensor, dict[str, Any]],
     targets: torch.Tensor,
     config: LVAELossConfig,
-    gaussian_likelihood: Optional[GaussianLikelihood],
-    noise_model_likelihood: Optional[NoiseModelLikelihood] = None,  # TODO: ugly
-) -> Optional[dict[str, torch.Tensor]]:
+    gaussian_likelihood: GaussianLikelihood | None,
+    noise_model_likelihood: NoiseModelLikelihood | None = None,  # TODO: ugly
+) -> dict[str, torch.Tensor] | None:
     """Loss function for muSplit.
 
     Parameters
@@ -351,9 +436,9 @@ def denoisplit_loss(
     model_outputs: tuple[torch.Tensor, dict[str, Any]],
     targets: torch.Tensor,
     config: LVAELossConfig,
-    gaussian_likelihood: Optional[GaussianLikelihood] = None,
-    noise_model_likelihood: Optional[NoiseModelLikelihood] = None,
-) -> Optional[dict[str, torch.Tensor]]:
+    gaussian_likelihood: GaussianLikelihood | None = None,
+    noise_model_likelihood: NoiseModelLikelihood | None = None,
+) -> dict[str, torch.Tensor] | None:
     """Loss function for DenoiSplit.
 
     Parameters
@@ -430,7 +515,7 @@ def denoisplit_musplit_loss(
     config: LVAELossConfig,
     gaussian_likelihood: GaussianLikelihood,
     noise_model_likelihood: NoiseModelLikelihood,
-) -> Optional[dict[str, torch.Tensor]]:
+) -> dict[str, torch.Tensor] | None:
     """Loss function for DenoiSplit.
 
     Parameters

careamics/lvae_training/dataset/__init__.py

@@ -1,4 +1,4 @@
-from .config import DatasetConfig
+from .config import MicroSplitDataConfig
 from .lc_dataset import LCMultiChDloader
 from .ms_dataset_ref import MultiChDloaderRef
 from .multich_dataset import MultiChDloader
@@ -7,14 +7,14 @@ from .multifile_dataset import MultiFileDset
 from .types import DataSplitType, DataType, TilingMode
 
 __all__ = [
-    "DatasetConfig",
-    "MultiChDloader",
+    "DataSplitType",
+    "DataType",
     "LCMultiChDloader",
-    "MultiFileDset",
-    "MultiCropDset",
-    "MultiChDloaderRef",
     "LCMultiChDloaderRef",
-    "DataType",
-    "DataSplitType",
+    "MicroSplitDataConfig",
+    "MultiChDloader",
+    "MultiChDloaderRef",
+    "MultiCropDset",
+    "MultiFileDset",
     "TilingMode",
 ]

careamics/lvae_training/dataset/config.py

@@ -1,4 +1,4 @@
-from typing import Any, Optional, Union
+from typing import Any, Union
 
 from pydantic import BaseModel, ConfigDict
 
@@ -6,70 +6,70 @@ from .types import DataSplitType, DataType, TilingMode
 
 
 # TODO: check if any bool logic can be removed
-class DatasetConfig(BaseModel):
-    model_config = ConfigDict(validate_assignment=True, extra="forbid")
+class MicroSplitDataConfig(BaseModel):
+    model_config = ConfigDict(validate_assignment=True, extra="allow")
 
-    data_type: Optional[DataType]
+    data_type: Union[DataType, str] | None  # TODO remove or refactor!!
     """Type of the dataset, should be one of DataType"""
 
-    depth3D: Optional[int] = 1
+    depth3D: int | None = 1
     """Number of slices in 3D. If data is 2D depth3D is equal to 1"""
 
-    datasplit_type: Optional[DataSplitType] = None
-    """Whether to return training, validation or test split, should be one of 
+    datasplit_type: DataSplitType | None = None
+    """Whether to return training, validation or test split, should be one of
     DataSplitType"""
 
-    num_channels: Optional[int] = 2
+    num_channels: int | None = 2
     """Number of channels in the input"""
 
     # TODO: remove ch*_fname parameters, should be parsed automatically from a name list
-    ch1_fname: Optional[str] = None
-    ch2_fname: Optional[str] = None
-    ch_input_fname: Optional[str] = None
+    ch1_fname: str | None = None
+    ch2_fname: str | None = None
+    ch_input_fname: str | None = None
 
-    input_is_sum: Optional[bool] = False
+    input_is_sum: bool | None = False
     """Whether the input is the sum or average of channels"""
 
-    input_idx: Optional[int] = None
+    input_idx: int | None = None
     """Index of the channel where the input is stored in the data"""
 
-    target_idx_list: Optional[list[int]] = None
+    target_idx_list: list[int] | None = None
     """Indices of the channels where the targets are stored in the data"""
 
     # TODO: where are there used?
-    start_alpha: Optional[Any] = None
-    end_alpha: Optional[Any] = None
+    start_alpha: Any | None = None
+    end_alpha: Any | None = None
 
     image_size: tuple  # TODO: revisit, new model_config uses tuple
     """Size of one patch of data"""
 
-    grid_size: Optional[Union[int, tuple[int, int, int]]] = None
+    grid_size: Union[int, tuple[int, int, int]] | None = None
     """Frame is divided into square grids of this size. A patch centered on a grid 
     having size `image_size` is returned. Grid size not used in training,
     used only during val / test, grid size controls the overlap of the patches"""
 
-    empty_patch_replacement_enabled: Optional[bool] = False
+    empty_patch_replacement_enabled: bool | None = False
     """Whether to replace the content of one of the channels
     with background with given probability"""
-    empty_patch_replacement_channel_idx: Optional[Any] = None
-    empty_patch_replacement_probab: Optional[Any] = None
-    empty_patch_max_val_threshold: Optional[Any] = None
+    empty_patch_replacement_channel_idx: Any | None = None
+    empty_patch_replacement_probab: Any | None = None
+    empty_patch_max_val_threshold: Any | None = None
 
-    uncorrelated_channels: Optional[bool] = False
-    """Replace the content in one of the channels with given probability to make 
+    uncorrelated_channels: bool | None = False
+    """Replace the content in one of the channels with given probability to make
     channel content 'uncorrelated'"""
-    uncorrelated_channel_probab: Optional[float] = 0.5
+    uncorrelated_channel_probab: float | None = 0.5
 
-    poisson_noise_factor: Optional[float] = -1
+    poisson_noise_factor: float | None = -1
     """The added poisson noise factor"""
 
-    synthetic_gaussian_scale: Optional[float] = 0.1
+    synthetic_gaussian_scale: float | None = 0.1
 
     # TODO: set to True in training code, recheck
-    input_has_dependant_noise: Optional[bool] = False
+    input_has_dependant_noise: bool | None = False
 
     # TODO: sometimes max_val differs between runs with fixed seeds with noise enabled
-    enable_gaussian_noise: Optional[bool] = False
+    enable_gaussian_noise: bool | None = False
     """Whether to enable gaussian noise"""
 
     # TODO: is this parameter used?
@@ -80,44 +80,56 @@ class DatasetConfig(BaseModel):
     deterministic_grid: Any = None
 
     # TODO: why is this not used?
-    enable_rotation_aug: Optional[bool] = False
+    enable_rotation_aug: bool | None = False
 
-    max_val: Optional[Union[float, tuple]] = None
-    """Maximum data in the dataset. Is calculated for train split, and should be 
+    max_val: Union[float, tuple] | None = None
+    """Maximum data in the dataset. Is calculated for train split, and should be
     externally set for val and test splits."""
 
     overlapping_padding_kwargs: Any = None
     """Parameters for np.pad method"""
 
     # TODO: remove this parameter, controls debug print
-    print_vars: Optional[bool] = False
+    print_vars: bool | None = False
 
     # Hard-coded parameters (used to be in the config file)
     normalized_input: bool = True
     """If this is set to true, then one mean and stdev is used
                 for both channels. Otherwise, two different mean and stdev are used."""
-    use_one_mu_std: Optional[bool] = True
+    use_one_mu_std: bool | None = True
 
     # TODO: is this parameter used?
-    train_aug_rotate: Optional[bool] = False
-    enable_random_cropping: Optional[bool] = True
+    train_aug_rotate: bool | None = False
+    enable_random_cropping: bool | None = True
 
-    multiscale_lowres_count: Optional[int] = None
+    multiscale_lowres_count: int | None = None
     """Number of LC scales"""
 
-    tiling_mode: Optional[TilingMode] = TilingMode.ShiftBoundary
+    tiling_mode: TilingMode | None = TilingMode.ShiftBoundary
 
-    target_separate_normalization: Optional[bool] = True
+    target_separate_normalization: bool | None = True
 
-    mode_3D: Optional[bool] = False
+    mode_3D: bool | None = False
     """If training in 3D mode or not"""
 
-    trainig_datausage_fraction: Optional[float] = 1.0
+    trainig_datausage_fraction: float | None = 1.0
 
-    validtarget_random_fraction: Optional[float] = None
+    validtarget_random_fraction: float | None = None
 
-    validation_datausage_fraction: Optional[float] = 1.0
+    validation_datausage_fraction: float | None = 1.0
 
-    random_flip_z_3D: Optional[bool] = False
+    random_flip_z_3D: bool | None = False
 
-    padding_kwargs: Optional[dict] = None
+    padding_kwargs: dict = {"mode": "reflect"}  # TODO remove !!
+
+    def __init__(self, **data):
+        # Convert string data_type to enum if needed
+        if "data_type" in data and isinstance(data["data_type"], str):
+            try:
+                data["data_type"] = DataType[data["data_type"]]
+            except KeyError:
+                # Keep original value to let validation handle the error
+                pass
+        super().__init__(**data)
+
+    # TODO add validators !

careamics/lvae_training/dataset/lc_dataset.py

@@ -2,23 +2,29 @@
 A place for Datasets and Dataloaders.
 """
 
-from typing import Tuple, Union, Callable
+import logging
+import math
+from pathlib import Path
+from typing import Any, Callable, Optional, Union
 
 import numpy as np
 from skimage.transform import resize
 
-from .config import DatasetConfig
+from .config import MicroSplitDataConfig
 from .multich_dataset import MultiChDloader
 
 
 class LCMultiChDloader(MultiChDloader):
+    """Multi-channel dataset loader for LC-style datasets."""
+
     def __init__(
         self,
-        data_config: DatasetConfig,
-        fpath: str,
-        load_data_fn: Callable,
-        val_fraction=None,
-        test_fraction=None,
+        data_config: MicroSplitDataConfig,
+        datapath: Union[str, Path],
+        load_data_fn: Optional[Callable] = None,
+        val_fraction: float = 0.1,
+        test_fraction: float = 0.1,
+        allow_generation: bool = False,
     ):
         self._padding_kwargs = (
             data_config.padding_kwargs  # mode=padding_mode, constant_values=constant_value
@@ -27,7 +33,7 @@ class LCMultiChDloader(MultiChDloader):
 
         super().__init__(
             data_config,
-            fpath,
+            datapath,
             load_data_fn=load_data_fn,
             val_fraction=val_fraction,
             test_fraction=test_fraction,
@@ -111,8 +117,8 @@ class LCMultiChDloader(MultiChDloader):
         return msg
 
     def _load_scaled_img(
-        self, scaled_index, index: Union[int, Tuple[int, int]]
-    ) -> Tuple[np.ndarray, np.ndarray]:
+        self, scaled_index, index: Union[int, tuple[int, int]]
+    ) -> tuple[np.ndarray, np.ndarray]:
         if isinstance(index, int):
             idx = index
         else:
@@ -131,7 +137,7 @@ class LCMultiChDloader(MultiChDloader):
             imgs = tuple([img + noise[0] * factor for img in imgs])
         return imgs
 
-    def _crop_img(self, img: np.ndarray, patch_start_loc: Tuple):
+    def _crop_img(self, img: np.ndarray, patch_start_loc: tuple):
         """
         Here, h_start, w_start could be negative. That simply means we need to pick the content from 0. So,
         the cropped image will be smaller than self._img_sz * self._img_sz
@@ -202,7 +208,7 @@ class LCMultiChDloader(MultiChDloader):
         )
         return output_img_tuples, cropped_noise_tuples
 
-    def __getitem__(self, index: Union[int, Tuple[int, int]]):
+    def __getitem__(self, index: Union[int, tuple[int, int]]):
         img_tuples, noise_tuples = self._get_img(index)
         if self._uncorrelated_channels:
             assert (

careamics/lvae_training/dataset/ms_dataset_ref.py

@@ -10,7 +10,7 @@ from typing import Callable, Union
 import numpy as np
 from skimage.transform import resize
 
-from .config import DatasetConfig
+from .config import MicroSplitDataConfig
 from .types import DataSplitType, TilingMode
 from .utils.empty_patch_fetcher import EmptyPatchFetcher
 from .utils.index_manager import GridIndexManagerRef
@@ -19,7 +19,7 @@ from .utils.index_manager import GridIndexManagerRef
 class MultiChDloaderRef:
     def __init__(
         self,
-        data_config: DatasetConfig,
+        data_config: MicroSplitDataConfig,
         fpath: str,
         load_data_fn: Callable,
         val_fraction: float = None,
@@ -171,8 +171,8 @@ class MultiChDloaderRef:
 
     def load_data(
         self,
-        data_config,
-        datasplit_type,
+        data_config: MicroSplitDataConfig,
+        datasplit_type: DataSplitType,
         load_data_fn: Callable,
         val_fraction=None,
         test_fraction=None,
@@ -813,7 +813,7 @@ class MultiChDloaderRef:
 class LCMultiChDloaderRef(MultiChDloaderRef):
     def __init__(
         self,
-        data_config: DatasetConfig,
+        data_config: MicroSplitDataConfig,
         fpath: str,
         load_data_fn: Callable,
         val_fraction=None,