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

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

@@ -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
 
@@ -269,6 +270,90 @@ 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,

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,

careamics/lvae_training/dataset/multich_dataset.py

@@ -2,29 +2,35 @@
 A place for Datasets and Dataloaders.
 """
 
-from typing import Tuple, Union, Callable
+from pathlib import Path
+from typing import Any, Callable, Optional, Union
 
 import numpy as np
+import torch
+from torch.utils.data import Dataset
 
 from .utils.empty_patch_fetcher import EmptyPatchFetcher
 from .utils.index_manager import GridIndexManager
 from .utils.index_switcher import IndexSwitcher
-from .config import DatasetConfig
+from .config import MicroSplitDataConfig
 from .types import DataSplitType, TilingMode
 
 
-class MultiChDloader:
+class MultiChDloader(Dataset):
+    """Multi-channel dataset loader."""
+
     def __init__(
         self,
-        data_config: DatasetConfig,
-        fpath: str,
-        load_data_fn: Callable,
-        val_fraction: float = None,
-        test_fraction: float = 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._data_type = data_config.data_type
-        self._fpath = fpath
+        self._fpath = datapath
         self._data = self._noise_data = None
         self.Z = 1
         self._5Ddata = False
@@ -395,7 +401,7 @@ class MultiChDloader:
         )
 
     def get_idx_manager_shapes(
-        self, patch_size: int, grid_size: Union[int, Tuple[int, int, int]]
+        self, patch_size: int, grid_size: Union[int, tuple[int, int, int]]
     ):
         numC = self._data.shape[-1]
         if self._5Ddata:
@@ -415,7 +421,7 @@ class MultiChDloader:
 
         return patch_shape, grid_shape
 
-    def set_img_sz(self, image_size, grid_size: Union[int, Tuple[int, int, int]]):
+    def set_img_sz(self, image_size, grid_size: Union[int, tuple[int, int, int]]):
         """
         If one wants to change the image size on the go, then this can be used.
         Args:
@@ -519,7 +525,7 @@ class MultiChDloader:
             },
         )
 
-    def _crop_img(self, img: np.ndarray, patch_start_loc: Tuple):
+    def _crop_img(self, img: np.ndarray, patch_start_loc: tuple):
         if self._tiling_mode in [TilingMode.TrimBoundary, TilingMode.ShiftBoundary]:
             # In training, this is used.
             # NOTE: It is my opinion that if I just use self._crop_img_with_padding, it will work perfectly fine.
@@ -600,7 +606,7 @@ class MultiChDloader:
         return new_img
 
     def _crop_flip_img(
-        self, img: np.ndarray, patch_start_loc: Tuple, h_flip: bool, w_flip: bool
+        self, img: np.ndarray, patch_start_loc: tuple, h_flip: bool, w_flip: bool
     ):
         new_img = self._crop_img(img, patch_start_loc)
         if h_flip:
@@ -611,8 +617,8 @@ class MultiChDloader:
         return new_img.astype(np.float32)
 
     def _load_img(
-        self, index: Union[int, Tuple[int, int]]
-    ) -> Tuple[np.ndarray, np.ndarray]:
+        self, index: Union[int, tuple[int, int]]
+    ) -> tuple[np.ndarray, np.ndarray]:
         """
         Returns the channels and also the respective noise channels.
         """
@@ -806,7 +812,7 @@ class MultiChDloader:
             w_start = 0
         return h_start, w_start
 
-    def _get_img(self, index: Union[int, Tuple[int, int]]):
+    def _get_img(self, index: Union[int, tuple[int, int]]):
         """
         Loads an image.
         Crops the image such that cropped image has content.
@@ -1056,8 +1062,8 @@ class MultiChDloader:
         return img_tuples, noise_tuples
 
     def __getitem__(
-        self, index: Union[int, Tuple[int, int]]
-    ) -> Tuple[np.ndarray, np.ndarray]:
+        self, index: Union[int, tuple[int, int]]
+    ) -> tuple[np.ndarray, np.ndarray]:
         # Vera: input can be both real microscopic image and two separate channels that are summed in the code
 
         if self._train_index_switcher is not None:

careamics/lvae_training/dataset/multifile_dataset.py

@@ -4,7 +4,7 @@ from typing import Callable, Union
 import numpy as np
 from numpy.typing import NDArray
 
-from .config import DatasetConfig
+from .config import MicroSplitDataConfig
 from .lc_dataset import LCMultiChDloader
 from .multich_dataset import MultiChDloader
 from .types import DataSplitType
@@ -82,7 +82,7 @@ class SingleFileLCDset(LCMultiChDloader):
     def __init__(
         self,
         preloaded_data: NDArray,
-        data_config: DatasetConfig,
+        data_config: MicroSplitDataConfig,
         fpath: str,
         load_data_fn: Callable,
         val_fraction=None,
@@ -106,7 +106,7 @@ class SingleFileLCDset(LCMultiChDloader):
 
     def load_data(
         self,
-        data_config: DatasetConfig,
+        data_config: MicroSplitDataConfig,
         datasplit_type: DataSplitType,
         load_data_fn: Callable,
         val_fraction=None,
@@ -124,7 +124,7 @@ class SingleFileDset(MultiChDloader):
     def __init__(
         self,
         preloaded_data: NDArray,
-        data_config: DatasetConfig,
+        data_config: MicroSplitDataConfig,
         fpath: str,
         load_data_fn: Callable,
         val_fraction=None,
@@ -148,7 +148,7 @@ class SingleFileDset(MultiChDloader):
 
     def load_data(
         self,
-        data_config: DatasetConfig,
+        data_config: MicroSplitDataConfig,
         datasplit_type: DataSplitType,
         load_data_fn: Callable[..., NDArray],
         val_fraction=None,
@@ -175,7 +175,7 @@ class MultiFileDset:
 
     def __init__(
         self,
-        data_config: DatasetConfig,
+        data_config: MicroSplitDataConfig,
         fpath: str,
         load_data_fn: Callable[..., Union[TwoChannelData, MultiChannelData]],
         val_fraction=None,

careamics/model_io/bmz_io.py

@@ -186,11 +186,15 @@ def export_to_bmz(
         )
 
         # test model description
-        test_kwargs = (
-            model_description.config.bioimageio.model_dump()
-            .get("test_kwargs", {})
-            .get("pytorch_state_dict", {})
-        )
+        test_kwargs = {}
+        if hasattr(model_description, "config") and isinstance(
+            model_description.config, dict
+        ):
+            bioimageio_config = model_description.config.get("bioimageio", {})
+            test_kwargs = bioimageio_config.get("test_kwargs", {}).get(
+                "pytorch_state_dict", {}
+            )
+
         summary: ValidationSummary = test_model(model_description, **test_kwargs)
         if summary.status == "failed":
             raise ValueError(f"Model description test failed: {summary}")