careamics 0.0.5__py3-none-any.whl → 0.0.7__py3-none-any.whl

careamics/lightning/callbacks/prediction_writer_callback/prediction_writer_callback.py

@@ -2,8 +2,9 @@
 
 from __future__ import annotations
 
+from collections.abc import Sequence
 from pathlib import Path
-from typing import Any, Optional, Sequence, Union
+from typing import Any, Optional, Union
 
 from pytorch_lightning import LightningModule, Trainer
 from pytorch_lightning.callbacks import BasePredictionWriter

careamics/lightning/callbacks/prediction_writer_callback/write_strategy.py

@@ -1,7 +1,8 @@
 """Module containing different strategies for writing predictions."""
 
+from collections.abc import Sequence
 from pathlib import Path
-from typing import Any, Optional, Protocol, Sequence, Union
+from typing import Any, Optional, Protocol, Union
 
 import numpy as np
 from numpy.typing import NDArray

careamics/lightning/callbacks/progress_bar_callback.py

@@ -1,11 +1,11 @@
 """Progressbar callback."""
 
 import sys
-from typing import Dict, Union
+from typing import Union
 
 from pytorch_lightning import LightningModule, Trainer
 from pytorch_lightning.callbacks import TQDMProgressBar
-from tqdm import tqdm
+from tqdm.auto import tqdm
 
 
 class ProgressBarCallback(TQDMProgressBar):
@@ -71,7 +71,7 @@ class ProgressBarCallback(TQDMProgressBar):
 
     def get_metrics(
         self, trainer: Trainer, pl_module: LightningModule
-    ) -> Dict[str, Union[int, str, float, Dict[str, float]]]:
+    ) -> dict[str, Union[int, str, float, dict[str, float]]]:
         """Override this to customize the metrics displayed in the progress bar.
 
         Parameters

careamics/lightning/lightning_module.py

@@ -6,7 +6,7 @@ import numpy as np
 import pytorch_lightning as L
 from torch import Tensor, nn
 
-from careamics.config import FCNAlgorithmConfig, VAEAlgorithmConfig
+from careamics.config import UNetBasedAlgorithm, VAEBasedAlgorithm
 from careamics.config.support import (
     SupportedAlgorithm,
     SupportedArchitecture,
@@ -34,6 +34,7 @@ from careamics.utils.torch_utils import get_optimizer, get_scheduler
 NoiseModel = Union[GaussianMixtureNoiseModel, MultiChannelNoiseModel]
 
 
+# TODO rename to UNetModule
 class FCNModule(L.LightningModule):
     """
     CAREamics Lightning module.
@@ -60,7 +61,7 @@ class FCNModule(L.LightningModule):
         Learning rate scheduler name.
     """
 
-    def __init__(self, algorithm_config: Union[FCNAlgorithmConfig, dict]) -> None:
+    def __init__(self, algorithm_config: Union[UNetBasedAlgorithm, dict]) -> None:
         """Lightning module for CAREamics.
 
         This class encapsulates the a PyTorch model along with the training, validation,
@@ -74,7 +75,9 @@ class FCNModule(L.LightningModule):
         super().__init__()
         # if loading from a checkpoint, AlgorithmModel needs to be instantiated
         if isinstance(algorithm_config, dict):
-            algorithm_config = FCNAlgorithmConfig(**algorithm_config)
+            algorithm_config = UNetBasedAlgorithm(
+                **algorithm_config
+            )  # TODO this needs to be updated using the algorithm-specific class
 
         # create model and loss function
         self.model: nn.Module = model_factory(algorithm_config.model)
@@ -266,7 +269,7 @@ class VAEModule(L.LightningModule):
         Learning rate scheduler name.
     """
 
-    def __init__(self, algorithm_config: Union[VAEAlgorithmConfig, dict]) -> None:
+    def __init__(self, algorithm_config: Union[VAEBasedAlgorithm, dict]) -> None:
         """Lightning module for CAREamics.
 
         This class encapsulates the a PyTorch model along with the training, validation,
@@ -280,7 +283,7 @@ class VAEModule(L.LightningModule):
         super().__init__()
         # if loading from a checkpoint, AlgorithmModel needs to be instantiated
         self.algorithm_config = (
-            VAEAlgorithmConfig(**algorithm_config)
+            VAEBasedAlgorithm(**algorithm_config)
             if isinstance(algorithm_config, dict)
             else algorithm_config
         )
@@ -656,9 +659,10 @@ def create_careamics_module(
     algorithm_configuration["model"] = model_configuration
 
     # call the parent init using an AlgorithmModel instance
+    # TODO broken by new configutations!
     algorithm_str = algorithm_configuration["algorithm"]
-    if algorithm_str in FCNAlgorithmConfig.get_compatible_algorithms():
-        return FCNModule(FCNAlgorithmConfig(**algorithm_configuration))
+    if algorithm_str in UNetBasedAlgorithm.get_compatible_algorithms():
+        return FCNModule(UNetBasedAlgorithm(**algorithm_configuration))
     else:
         raise NotImplementedError(
             f"Model {algorithm_str} is not implemented or unknown."

careamics/lightning/train_data_module.py

@@ -2,14 +2,13 @@
 
 from pathlib import Path
 from typing import Any, Callable, Literal, Optional, Union
-from warnings import warn
 
 import numpy as np
 import pytorch_lightning as L
 from numpy.typing import NDArray
 from torch.utils.data import DataLoader, IterableDataset
 
-from careamics.config import DataConfig
+from careamics.config.data import DataConfig, GeneralDataConfig, N2VDataConfig
 from careamics.config.support import SupportedData
 from careamics.config.transformations import TransformModel
 from careamics.dataset.dataset_utils import (
@@ -119,7 +118,7 @@ class TrainDataModule(L.LightningDataModule):
 
     def __init__(
         self,
-        data_config: DataConfig,
+        data_config: GeneralDataConfig,
         train_data: Union[Path, str, NDArray],
         val_data: Optional[Union[Path, str, NDArray]] = None,
         train_data_target: Optional[Union[Path, str, NDArray]] = None,
@@ -219,7 +218,7 @@ class TrainDataModule(L.LightningDataModule):
             )
 
         # configuration
-        self.data_config: DataConfig = data_config
+        self.data_config: GeneralDataConfig = data_config
         self.data_type: str = data_config.data_type
         self.batch_size: int = data_config.batch_size
         self.use_in_memory: bool = use_in_memory
@@ -261,11 +260,6 @@ class TrainDataModule(L.LightningDataModule):
 
         self.extension_filter: str = extension_filter
 
-        # Pytorch dataloader parameters
-        self.dataloader_params: dict[str, Any] = (
-            data_config.dataloader_params if data_config.dataloader_params else {}
-        )
-
     def prepare_data(self) -> None:
         """
         Hook used to prepare the data before calling `setup`.
@@ -447,21 +441,17 @@ class TrainDataModule(L.LightningDataModule):
         Any
             Training dataloader.
         """
-        # check because iterable dataset cannot be shuffled
-        if not isinstance(self.train_dataset, IterableDataset):
-            if ("shuffle" in self.dataloader_params) and (
-                not self.dataloader_params["shuffle"]
-            ):
-                warn(
-                    "Dataloader parameters include `shuffle=False`, this will be "
-                    "passed to the training dataloader and may result in bad results.",
-                    stacklevel=1,
-                )
-            else:
-                self.dataloader_params["shuffle"] = True
+        train_dataloader_params = self.data_config.train_dataloader_params.copy()
+
+        # NOTE: When next-gen datasets are completed this can be removed
+        # iterable dataset cannot be shuffled
+        if isinstance(self.train_dataset, IterableDataset):
+            del train_dataloader_params["shuffle"]
 
         return DataLoader(
-            self.train_dataset, batch_size=self.batch_size, **self.dataloader_params
+            self.train_dataset,
+            batch_size=self.batch_size,
+            **train_dataloader_params,
         )
 
     def val_dataloader(self) -> Any:
@@ -476,6 +466,7 @@ class TrainDataModule(L.LightningDataModule):
         return DataLoader(
             self.val_dataset,
             batch_size=self.batch_size,
+            **self.data_config.val_dataloader_params,
         )
 
 
@@ -502,12 +493,23 @@ def create_train_datamodule(
     """Create a TrainDataModule.
 
     This function is used to explicitly pass the parameters usually contained in a
-    `data_model` configuration to a TrainDataModule.
+    `GenericDataConfig` to a TrainDataModule.
 
     Since the lightning datamodule has no access to the model, make sure that the
     parameters passed to the datamodule are consistent with the model's requirements and
     are coherent.
 
+    By default, the train DataModule will be set for Noise2Void if no target data is
+    provided. That means that it will add a `N2VManipulateModel` transformation to the
+    list of augmentations. The default augmentations are XY flip, XY rotation, and N2V
+    pixel manipulation. If you pass a training target data, the default behaviour is to
+    train a supervised model. It will use the default XY flip and rotation
+    augmentations.
+
+    To use a different set of transformations, you can pass a list of transforms to
+    `transforms`. Note that if you intend to use Noise2Void, you should add
+    `N2VManipulateModel` as the last transform in the list of transformations.
+
     The data module can be used with Path, str or numpy arrays. In the case of
     numpy arrays, it loads and computes all the patches in memory. For Path and str
     inputs, it calculates the total file size and estimate whether it can fit in
@@ -518,11 +520,6 @@ def create_train_datamodule(
     To use array data, set `data_type` to `array` and pass a numpy array to
     `train_data`.
 
-    In particular, N2V requires a specific transformation (N2V manipulates), which is
-    not compatible with supervised training. The default transformations applied to the
-    training patches are defined in `careamics.config.data_model`. To use different
-    transformations, pass a list of transforms. See examples for more details.
-
     By default, CAREamics only supports types defined in
     `careamics.config.support.SupportedData`. To read custom data types, you can set
     `data_type` to `custom` and provide a function that returns a numpy array from a
@@ -627,12 +624,12 @@ def create_train_datamodule(
     transforms:
     >>> import numpy as np
     >>> from careamics.lightning import create_train_datamodule
+    >>> from careamics.config.transformations import XYFlipModel, N2VManipulateModel
     >>> from careamics.config.support import SupportedTransform
     >>> my_array = np.arange(256).reshape(16, 16)
     >>> my_transforms = [
-    ...     {
-    ...         "name": SupportedTransform.XY_FLIP.value,
-    ...     }
+    ...     XYFlipModel(flip_y=False),
+    ...     N2VManipulateModel()
     ... ]
     >>> data_module = create_train_datamodule(
     ...     train_data=my_array,
@@ -659,21 +656,15 @@ def create_train_datamodule(
     if transforms is not None:
         data_dict["transforms"] = transforms
 
-    # validate configuration
-    data_config = DataConfig(**data_dict)
+    # TODO not compatible with HDN, consider adding an argument for n2v/hdn
+    if train_target_data is None:
+        data_config: GeneralDataConfig = N2VDataConfig(**data_dict)
+        assert isinstance(data_config, N2VDataConfig)
 
-    # N2V specific checks, N2V, structN2V, and transforms
-    if data_config.has_n2v_manipulate():
-        # there is not target, n2v2 and structN2V can be changed
-        if train_target_data is None:
-            data_config.set_N2V2(use_n2v2)
-            data_config.set_structN2V_mask(struct_n2v_axis, struct_n2v_span)
-        else:
-            raise ValueError(
-                "Cannot have both supervised training (target data) and "
-                "N2V manipulation in the transforms. Pass a list of transforms "
-                "that is compatible with your supervised training."
-            )
+        data_config.set_n2v2(use_n2v2)
+        data_config.set_structN2V_mask(struct_n2v_axis, struct_n2v_span)
+    else:
+        data_config = DataConfig(**data_dict)
 
     # sanity check on the dataloader parameters
     if "batch_size" in dataloader_params:

careamics/losses/__init__.py

@@ -1,13 +1,13 @@
 """Losses module."""
 
 __all__ = [
+    "denoisplit_loss",
+    "denoisplit_musplit_loss",
     "loss_factory",
     "mae_loss",
     "mse_loss",
-    "n2v_loss",
-    "denoisplit_loss",
     "musplit_loss",
-    "denoisplit_musplit_loss",
+    "n2v_loss",
 ]
 
 from .fcn.losses import mae_loss, mse_loss, n2v_loss

careamics/lvae_training/calibration.py

@@ -1,4 +1,4 @@
-from typing import Union
+from typing import Union, Optional
 
 import numpy as np
 import torch
@@ -34,9 +34,6 @@ class Calibration:
         self._bins = num_bins
         self._bin_boundaries = None
 
-    def logvar_to_std(self, logvar: np.ndarray) -> np.ndarray:
-        return np.exp(logvar / 2)
-
     def compute_bin_boundaries(self, predict_std: np.ndarray) -> np.ndarray:
         """Compute the bin boundaries for `num_bins` bins and predicted std values."""
         min_std = np.min(predict_std)
@@ -104,65 +101,75 @@ class Calibration:
                 )
                 rmse_stderr = np.sqrt(stderr) if stderr is not None else None
 
-                bin_var = np.mean((std_ch[bin_mask] ** 2))
+                bin_var = np.mean(std_ch[bin_mask] ** 2)
                 stats_dict[ch_idx]["rmse"].append(bin_error)
                 stats_dict[ch_idx]["rmse_err"].append(rmse_stderr)
                 stats_dict[ch_idx]["rmv"].append(np.sqrt(bin_var))
                 stats_dict[ch_idx]["bin_count"].append(bin_size)
+        self.stats_dict = stats_dict
         return stats_dict
 
+    def get_calibrated_factor_for_stdev(
+        self,
+        pred: Optional[np.ndarray] = None,
+        pred_std: Optional[np.ndarray] = None,
+        target: Optional[np.ndarray] = None,
+        q_s: float = 0.00001,
+        q_e: float = 0.99999,
+    ) -> dict[str, float]:
+        """Calibrate the uncertainty by multiplying the predicted std with a scalar.
 
-def get_calibrated_factor_for_stdev(
-    pred: Union[np.ndarray, torch.Tensor],
-    pred_std: Union[np.ndarray, torch.Tensor],
-    target: Union[np.ndarray, torch.Tensor],
-    q_s: float = 0.00001,
-    q_e: float = 0.99999,
-    num_bins: int = 30,
-) -> dict[str, float]:
-    """Calibrate the uncertainty by multiplying the predicted std with a scalar.
-
-    Parameters
-    ----------
-    pred : Union[np.ndarray, torch.Tensor]
-        Predicted image, shape (n, h, w, c).
-    pred_std : Union[np.ndarray, torch.Tensor]
-        Predicted std, shape (n, h, w, c).
-    target : Union[np.ndarray, torch.Tensor]
-        Target image, shape (n, h, w, c).
-    q_s : float, optional
-        Start quantile, by default 0.00001.
-    q_e : float, optional
-        End quantile, by default 0.99999.
-    num_bins : int, optional
-        Number of bins to use for calibration, by default 30.
-
-    Returns
-    -------
-    dict[str, float]
-        Calibrated factor for each channel (slope + intercept).
-    """
-    calib = Calibration(num_bins=num_bins)
-    stats_dict = calib.compute_stats(pred, pred_std, target)
-    outputs = {}
-    for ch_idx in stats_dict.keys():
-        y = stats_dict[ch_idx]["rmse"]
-        x = stats_dict[ch_idx]["rmv"]
-        count = stats_dict[ch_idx]["bin_count"]
-
-        first_idx = get_first_index(count, q_s)
-        last_idx = get_last_index(count, q_e)
-        x = x[first_idx:-last_idx]
-        y = y[first_idx:-last_idx]
-        slope, intercept, *_ = stats.linregress(x, y)
-        output = {"scalar": slope, "offset": intercept}
-        outputs[ch_idx] = output
-    return outputs
+        Parameters
+        ----------
+        stats_dict : dict[int, dict[str, Union[np.ndarray, list]]]
+            Dictionary containing the stats for each channel.
+        q_s : float, optional
+            Start quantile, by default 0.00001.
+        q_e : float, optional
+            End quantile, by default 0.99999.
+
+        Returns
+        -------
+        dict[str, float]
+            Calibrated factor for each channel (slope + intercept).
+        """
+        if not hasattr(self, "stats_dict"):
+            print("No stats found. Computing stats...")
+            if any(v is None for v in [pred, pred_std, target]):
+                raise ValueError("pred, pred_std, and target must be provided.")
+            self.stats_dict = self.compute_stats(
+                pred=pred, pred_std=pred_std, target=target
+            )
+        outputs = {}
+        for ch_idx in self.stats_dict.keys():
+            y = self.stats_dict[ch_idx]["rmse"]
+            x = self.stats_dict[ch_idx]["rmv"]
+            count = self.stats_dict[ch_idx]["bin_count"]
+
+            first_idx = get_first_index(count, q_s)
+            last_idx = get_last_index(count, q_e)
+            x = x[first_idx:-last_idx]
+            y = y[first_idx:-last_idx]
+            slope, intercept, *_ = stats.linregress(x, y)
+            output = {"scalar": slope, "offset": intercept}
+            outputs[ch_idx] = output
+        factors = self.get_factors_array(factors_dict=outputs)
+        return outputs, factors
+
+    def get_factors_array(self, factors_dict: list[dict]):
+        """Get the calibration factors as a numpy array."""
+        calib_scalar = [factors_dict[i]["scalar"] for i in range(len(factors_dict))]
+        calib_scalar = np.array(calib_scalar).reshape(1, 1, 1, -1)
+        calib_offset = [
+            factors_dict[i].get("offset", 0.0) for i in range(len(factors_dict))
+        ]
+        calib_offset = np.array(calib_offset).reshape(1, 1, 1, -1)
+        return {"scalar": calib_scalar, "offset": calib_offset}
 
 
 def plot_calibration(ax, calibration_stats):
-    first_idx = get_first_index(calibration_stats[0]["bin_count"], 0.001)
-    last_idx = get_last_index(calibration_stats[0]["bin_count"], 0.999)
+    first_idx = get_first_index(calibration_stats[0]["bin_count"], 0.0001)
+    last_idx = get_last_index(calibration_stats[0]["bin_count"], 0.9999)
     ax.plot(
         calibration_stats[0]["rmv"][first_idx:-last_idx],
         calibration_stats[0]["rmse"][first_idx:-last_idx],
@@ -170,15 +177,15 @@ def plot_calibration(ax, calibration_stats):
         label=r"$\hat{C}_0$: Ch1",
     )
 
-    first_idx = get_first_index(calibration_stats[1]["bin_count"], 0.001)
-    last_idx = get_last_index(calibration_stats[1]["bin_count"], 0.999)
+    first_idx = get_first_index(calibration_stats[1]["bin_count"], 0.0001)
+    last_idx = get_last_index(calibration_stats[1]["bin_count"], 0.9999)
     ax.plot(
         calibration_stats[1]["rmv"][first_idx:-last_idx],
         calibration_stats[1]["rmse"][first_idx:-last_idx],
         "o",
-        label=r"$\hat{C}_1: : Ch2$",
+        label=r"$\hat{C}_1$: Ch2",
     )
-
+    # TODO add multichannel
     ax.set_xlabel("RMV")
     ax.set_ylabel("RMSE")
     ax.legend()

careamics/lvae_training/dataset/lc_dataset.py

@@ -97,7 +97,8 @@ class LCMultiChDloader(MultiChDloader):
             ]
 
         self.N = len(t_list)
-        self.set_img_sz(self._img_sz, self._grid_sz)
+        # TODO where tf is self._img_sz defined?
+        self.set_img_sz([self._img_sz, self._img_sz], self._grid_sz)
         print(
             f"[{self.__class__.__name__}] Data reduced. New data shape: {self._data.shape}"
         )

careamics/lvae_training/dataset/multich_dataset.py

@@ -359,8 +359,8 @@ class MultiChDloader:
             self._noise_data = self._noise_data[
                 t_list, h_start:h_end, w_start:w_end, :
             ].copy()
-
-        self.set_img_sz(self._img_sz, self._grid_sz)
+        # TODO where tf is self._img_sz defined?
+        self.set_img_sz([self._img_sz, self._img_sz], self._grid_sz)
         print(
             f"[{self.__class__.__name__}] Data reduced. New data shape: {self._data.shape}"
         )

careamics/lvae_training/dataset/types.py

@@ -3,7 +3,7 @@ from enum import Enum
 
 class DataType(Enum):
     Elisa3DData = 0
-    NicolaData = 1
+    HTLIF24Data = 1
     Pavia3SeqData = 2
     TavernaSox2GolgiV2 = 3
     Dao3ChannelWithInput = 4