careamics 0.0.3__py3-none-any.whl → 0.0.4__py3-none-any.whl

careamics/models/lvae/likelihoods.py

@@ -7,6 +7,7 @@ from __future__ import annotations
 import math
 from typing import Literal, Union, TYPE_CHECKING, Any, Optional
 
+import numpy as np
 import torch
 from torch import nn
 
@@ -287,30 +288,37 @@ class NoiseModelLikelihood(LikelihoodModule):
 
     def __init__(
         self,
-        data_mean: torch.Tensor,
-        data_std: torch.Tensor,
-        noiseModel: NoiseModel,  # TODO: check the type -> couldn't manage due to circular imports...
+        data_mean: Union[np.ndarray, torch.Tensor],
+        data_std: Union[np.ndarray, torch.Tensor],
+        noiseModel: NoiseModel,
     ):
         """Constructor.
 
         Parameters
         ----------
-        data_mean: torch.Tensor
+        data_mean: Union[np.ndarray, torch.Tensor]
             The mean of the data, used to unnormalize data for noise model evaluation.
-        data_std: torch.Tensor
+        data_std: Union[np.ndarray, torch.Tensor]
             The standard deviation of the data, used to unnormalize data for noise
             model evaluation.
         noiseModel: NoiseModel
             The noise model instance used to compute the likelihood.
         """
         super().__init__()
-        self.data_mean = data_mean
-        self.data_std = data_std
+        self.data_mean = torch.Tensor(data_mean)
+        self.data_std = torch.Tensor(data_std)
         self.noiseModel = noiseModel
 
-    def set_params_to_same_device_as(
+    def _set_params_to_same_device_as(
         self, correct_device_tensor: torch.Tensor
-    ) -> None:  # TODO: needed?
+    ) -> None:
+        """Set the parameters to the same device as the input tensor.
+
+        Parameters
+        ----------
+        correct_device_tensor: torch.Tensor
+            The tensor whose device is used to set the parameters.
+        """
         if self.data_mean.device != correct_device_tensor.device:
             self.data_mean = self.data_mean.to(correct_device_tensor.device)
             self.data_std = self.data_std.to(correct_device_tensor.device)
@@ -355,6 +363,7 @@ class NoiseModelLikelihood(LikelihoodModule):
         torch.Tensor
             The log-likelihood tensor. Shape is (B, C, [Z], Y, X).
         """
+        self._set_params_to_same_device_as(x)
         predicted_s_denormalized = params["mean"] * self.data_std + self.data_mean
         x_denormalized = x * self.data_std + self.data_mean
         likelihoods = self.noiseModel.likelihood(

careamics/models/lvae/lvae.py

@@ -38,7 +38,6 @@ class LadderVAE(nn.Module):
         decoder_dropout: float,
         nonlinearity: str,
         predict_logvar: bool,
-        enable_noise_model: bool,
         analytical_kl: bool,
     ):
         """
@@ -62,15 +61,12 @@ class LadderVAE(nn.Module):
         self.decoder_dropout = decoder_dropout
         self.nonlin = nonlinearity
         self.predict_logvar = predict_logvar
-        self.enable_noise_model = enable_noise_model
-
         self.analytical_kl = analytical_kl
         # -------------------------------------------------------
 
         # -------------------------------------------------------
         # Model attributes -> Hardcoded
         self.model_type = ModelType.LadderVae  # TODO remove !
-        self.model_type = ModelType.LadderVae  # TODO remove !
         self.encoder_blocks_per_layer = 1
         self.decoder_blocks_per_layer = 1
         self.bottomup_batchnorm = True
@@ -94,13 +90,6 @@ class LadderVAE(nn.Module):
         self._stochastic_use_naive_exponential = False
         self._enable_topdown_normalize_factor = True
 
-        # Noise model attributes -> Hardcoded
-        self.noise_model_type = "gmm"
-        self.denoise_channel = (
-            "input"  # 4 values for denoise_channel {'Ch1', 'Ch2', 'input','all'}
-        )
-        self.noise_model_learnable = False
-
         # Attributes that handle LC -> Hardcoded
         self.enable_multiscale = (
             self._multiscale_count is not None and self._multiscale_count > 1
@@ -806,11 +795,18 @@ class LadderVAE(nn.Module):
 
     #     return samples
 
-    # def reset_for_different_output_size(self, output_size):
-    #     for i in range(self.n_layers):
-    #         sz = output_size // 2**(1 + i)
-    #         self.bottom_up_layers[i].output_expected_shape = (sz, sz)
-    #         self.top_down_layers[i].latent_shape = (output_size, output_size)
+    def reset_for_different_output_size(self, output_size: int) -> None:
+        """Reset shape of output and latent tensors for different output size.
+
+        Used during evaluation to reset expected shapes of tensors when
+        input/output shape changes.
+        For instance, it is needed when the model was trained on, say, 64x64 sized
+        patches, but prediction is done on 128x128 patches.
+        """
+        for i in range(self.n_layers):
+            sz = output_size // 2 ** (1 + i)
+            self.bottom_up_layers[i].output_expected_shape = (sz, sz)
+            self.top_down_layers[i].latent_shape = (output_size, output_size)
 
     def pad_input(self, x):
         """

careamics/models/lvae/noise_models.py

@@ -76,7 +76,7 @@ def train_gm_noise_model(
     # TODO any training params ? Different channels ?
     noise_model = GaussianMixtureNoiseModel(model_config)
     # TODO revisit config unpacking
-    noise_model.train_noise_model(noise_model.signal, noise_model.observation)
+    noise_model.train_noise_model(model_config.signal, model_config.observation)
     return noise_model
 
 

careamics/transforms/compose.py

@@ -1,13 +1,14 @@
 """A class chaining transforms together."""
 
-from typing import Dict, List, Optional, Tuple, cast
+from typing import Dict, List, Optional, Tuple, Union, cast
 
-import numpy as np
+from numpy.typing import NDArray
 
-from careamics.config.data_model import TRANSFORMS_UNION
+from careamics.config.transformations import TransformModel
 
 from .n2v_manipulate import N2VManipulate
 from .normalize import Normalize
+from .transform import Transform
 from .xy_flip import XYFlip
 from .xy_random_rotate90 import XYRandomRotate90
 
@@ -36,7 +37,7 @@ class Compose:
 
     Parameters
     ----------
-    transform_list : List[TRANSFORMS_UNION]
+    transform_list : List[TransformModel]
         A list of dictionaries where each dictionary contains the name of a
         transform and its parameters.
 
@@ -46,26 +47,27 @@ class Compose:
         A callable that applies the transforms to the input data.
     """
 
-    def __init__(self, transform_list: List[TRANSFORMS_UNION]) -> None:
+    def __init__(self, transform_list: List[TransformModel]) -> None:
         """Instantiate a Compose object.
 
         Parameters
         ----------
-        transform_list : List[TRANSFORMS_UNION]
+        transform_list : List[TransformModel]
             A list of dictionaries where each dictionary contains the name of a
             transform and its parameters.
         """
         # retrieve all available transforms
-        all_transforms = get_all_transforms()
+        # TODO: correctly type hint get_all_transforms function output
+        all_transforms: dict[str, type[Transform]] = get_all_transforms()
 
         # instantiate all transforms
-        self.transforms = [
+        self.transforms: list[Transform] = [
             all_transforms[t.name](**t.model_dump()) for t in transform_list
         ]
 
     def _chain_transforms(
-        self, patch: np.ndarray, target: Optional[np.ndarray]
-    ) -> Tuple[np.ndarray, Optional[np.ndarray]]:
+        self, patch: NDArray, target: Optional[NDArray]
+    ) -> Tuple[Optional[NDArray], ...]:
         """Chain transforms on the input data.
 
         Parameters
@@ -80,16 +82,56 @@ class Compose:
         Tuple[np.ndarray, Optional[np.ndarray]]
             The output of the transformations.
         """
-        params = (patch, target)
+        params: Union[
+            tuple[NDArray, Optional[NDArray]],
+            tuple[NDArray, NDArray, NDArray],  # N2VManiupulate output
+        ] = (patch, target)
 
         for t in self.transforms:
-            params = t(*params)
+            # N2VManipulate returns tuple of 3 arrays
+            #   - Other transoforms return tuple of (patch, target, additional_arrays)
+            if isinstance(t, N2VManipulate):
+                patch, *_ = params
+                params = t(patch=patch)
+            else:
+                *params, _ = t(*params)  # ignore additional_arrays dict
 
         return params
 
+    def _chain_transforms_additional_arrays(
+        self,
+        patch: NDArray,
+        target: Optional[NDArray],
+        **additional_arrays: NDArray,
+    ) -> Tuple[NDArray, Optional[NDArray], dict[str, NDArray]]:
+        """Chain transforms on the input data, with additional arrays.
+
+        Parameters
+        ----------
+        patch : np.ndarray
+            Input data.
+        target : Optional[np.ndarray]
+            Target data, by default None.
+        **additional_arrays : NDArray
+            Additional arrays that will be transformed identically to `patch` and
+            `target`.
+
+        Returns
+        -------
+        Tuple[np.ndarray, Optional[np.ndarray]]
+            The output of the transformations.
+        """
+        params = {"patch": patch, "target": target, **additional_arrays}
+
+        for t in self.transforms:
+            patch, target, additional_arrays = t(**params)
+            params = {"patch": patch, "target": target, **additional_arrays}
+
+        return patch, target, additional_arrays
+
     def __call__(
-        self, patch: np.ndarray, target: Optional[np.ndarray] = None
-    ) -> Tuple[np.ndarray, ...]:
+        self, patch: NDArray, target: Optional[NDArray] = None
+    ) -> Tuple[NDArray, ...]:
         """Apply the transforms to the input data.
 
         Parameters
@@ -104,4 +146,37 @@ class Compose:
         Tuple[np.ndarray, ...]
             The output of the transformations.
         """
-        return cast(Tuple[np.ndarray, ...], self._chain_transforms(patch, target))
+        # TODO: solve casting Compose.__call__ ouput
+        return cast(Tuple[NDArray, ...], self._chain_transforms(patch, target))
+
+    def transform_with_additional_arrays(
+        self,
+        patch: NDArray,
+        target: Optional[NDArray] = None,
+        **additional_arrays: NDArray,
+    ) -> tuple[NDArray, Optional[NDArray], dict[str, NDArray]]:
+        """Apply the transforms to the input data, including additional arrays.
+
+        Parameters
+        ----------
+        patch : np.ndarray
+            The input data.
+        target : Optional[np.ndarray], optional
+            Target data, by default None.
+        **additional_arrays : NDArray
+            Additional arrays that will be transformed identically to `patch` and
+            `target`.
+
+        Returns
+        -------
+        NDArray
+            The transformed patch.
+        NDArray | None
+            The transformed target.
+        dict of {str, NDArray}
+            Transformed additional arrays. Keys correspond to the keyword argument
+            names.
+        """
+        return self._chain_transforms_additional_arrays(
+            patch, target, **additional_arrays
+        )

careamics/transforms/n2v_manipulate.py

@@ -3,6 +3,7 @@
 from typing import Any, Literal, Optional, Tuple
 
 import numpy as np
+from numpy.typing import NDArray
 
 from careamics.config.support import SupportedPixelManipulation, SupportedStructAxis
 from careamics.transforms.transform import Transform
@@ -98,8 +99,8 @@ class N2VManipulate(Transform):
         self.rng = np.random.default_rng(seed=seed)
 
     def __call__(
-        self, patch: np.ndarray, *args: Any, **kwargs: Any
-    ) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
+        self, patch: NDArray, *args: Any, **kwargs: Any
+    ) -> Tuple[NDArray, NDArray, NDArray]:
         """Apply the transform to the image.
 
         Parameters
@@ -142,5 +143,8 @@ class N2VManipulate(Transform):
         else:
             raise ValueError(f"Unknown masking strategy ({self.strategy}).")
 
+        # TODO: Output does not match other transforms, how to resolve?
+        #     - Don't include in Compose and apply after if algorithm is N2V?
+        #     - or just don't return patch? but then mask is in the target position
         # TODO why return patch?
         return masked, patch, mask

careamics/transforms/normalize.py

@@ -90,8 +90,11 @@ class Normalize(Transform):
         self.eps = 1e-6
 
     def __call__(
-        self, patch: np.ndarray, target: Optional[NDArray] = None
-    ) -> tuple[NDArray, Optional[NDArray]]:
+        self,
+        patch: np.ndarray,
+        target: Optional[NDArray] = None,
+        **additional_arrays: NDArray,
+    ) -> tuple[NDArray, Optional[NDArray], dict[str, NDArray]]:
         """Apply the transform to the source patch and the target (optional).
 
         Parameters
@@ -100,6 +103,9 @@ class Normalize(Transform):
             Patch, 2D or 3D, shape C(Z)YX.
         target : NDArray, optional
             Target for the patch, by default None.
+        **additional_arrays : NDArray
+            Additional arrays that will be transformed identically to `patch` and
+            `target`.
 
         Returns
         -------
@@ -111,6 +117,11 @@ class Normalize(Transform):
                 f"Number of means (got a list of size {len(self.image_means)}) and "
                 f"number of channels (got shape {patch.shape} for C(Z)YX) do not match."
             )
+        if len(additional_arrays) != 0:
+            raise NotImplementedError(
+                "Transforming additional arrays is currently not supported for "
+                "`Normalize`."
+            )
 
         # reshape mean and std and apply the normalization to the patch
         means = _reshape_stats(self.image_means, patch.ndim)
@@ -129,7 +140,7 @@ class Normalize(Transform):
         else:
             norm_target = None
 
-        return norm_patch, norm_target
+        return norm_patch, norm_target, additional_arrays
 
     def _apply(self, patch: NDArray, mean: NDArray, std: NDArray) -> NDArray:
         """

careamics/transforms/xy_flip.py

@@ -1,8 +1,9 @@
 """XY flip transform."""
 
-from typing import Optional, Tuple
+from typing import Optional
 
 import numpy as np
+from numpy.typing import NDArray
 
 from careamics.transforms.transform import Transform
 
@@ -78,8 +79,11 @@ class XYFlip(Transform):
         self.rng = np.random.default_rng(seed=seed)
 
     def __call__(
-        self, patch: np.ndarray, target: Optional[np.ndarray] = None
-    ) -> Tuple[np.ndarray, Optional[np.ndarray]]:
+        self,
+        patch: NDArray,
+        target: Optional[NDArray] = None,
+        **additional_arrays: NDArray,
+    ) -> tuple[NDArray, Optional[NDArray], dict[str, NDArray]]:
         """Apply the transform to the source patch and the target (optional).
 
         Parameters
@@ -88,6 +92,9 @@ class XYFlip(Transform):
             Patch, 2D or 3D, shape C(Z)YX.
         target : Optional[np.ndarray], optional
             Target for the patch, by default None.
+        **additional_arrays : NDArray
+            Additional arrays that will be transformed identically to `patch` and
+            `target`.
 
         Returns
         -------
@@ -95,17 +102,20 @@ class XYFlip(Transform):
             Transformed patch and target.
         """
         if self.rng.random() > self.p:
-            return patch, target
+            return patch, target, additional_arrays
 
         # choose an axis to flip
         axis = self.rng.choice(self.axis_indices)
 
         patch_transformed = self._apply(patch, axis)
         target_transformed = self._apply(target, axis) if target is not None else None
+        additional_transformed = {
+            key: self._apply(array, axis) for key, array in additional_arrays.items()
+        }
 
-        return patch_transformed, target_transformed
+        return patch_transformed, target_transformed, additional_transformed
 
-    def _apply(self, patch: np.ndarray, axis: int) -> np.ndarray:
+    def _apply(self, patch: NDArray, axis: int) -> NDArray:
         """Apply the transform to the image.
 
         Parameters

careamics/transforms/xy_random_rotate90.py

@@ -3,6 +3,7 @@
 from typing import Optional, Tuple
 
 import numpy as np
+from numpy.typing import NDArray
 
 from careamics.transforms.transform import Transform
 
@@ -49,8 +50,11 @@ class XYRandomRotate90(Transform):
         self.rng = np.random.default_rng(seed=seed)
 
     def __call__(
-        self, patch: np.ndarray, target: Optional[np.ndarray] = None
-    ) -> Tuple[np.ndarray, Optional[np.ndarray]]:
+        self,
+        patch: NDArray,
+        target: Optional[NDArray] = None,
+        **additional_arrays: NDArray,
+    ) -> tuple[NDArray, Optional[NDArray], dict[str, NDArray]]:
         """Apply the transform to the source patch and the target (optional).
 
         Parameters
@@ -59,6 +63,9 @@ class XYRandomRotate90(Transform):
             Patch, 2D or 3D, shape C(Z)YX.
         target : Optional[np.ndarray], optional
             Target for the patch, by default None.
+        **additional_arrays : NDArray
+            Additional arrays that will be transformed identically to `patch` and
+            `target`.
 
         Returns
         -------
@@ -66,7 +73,7 @@ class XYRandomRotate90(Transform):
             Transformed patch and target.
         """
         if self.rng.random() > self.p:
-            return patch, target
+            return patch, target, additional_arrays
 
         # number of rotations
         n_rot = self.rng.integers(1, 4)
@@ -76,12 +83,14 @@ class XYRandomRotate90(Transform):
         target_transformed = (
             self._apply(target, n_rot, axes) if target is not None else None
         )
+        additional_transformed = {
+            key: self._apply(array, n_rot, axes)
+            for key, array in additional_arrays.items()
+        }
 
-        return patch_transformed, target_transformed
+        return patch_transformed, target_transformed, additional_transformed
 
-    def _apply(
-        self, patch: np.ndarray, n_rot: int, axes: Tuple[int, int]
-    ) -> np.ndarray:
+    def _apply(self, patch: NDArray, n_rot: int, axes: Tuple[int, int]) -> NDArray:
         """Apply the transform to the image.
 
         Parameters