careamics 0.1.0rc5__py3-none-any.whl → 0.1.0rc7__py3-none-any.whl

careamics/transforms/normalize.py

@@ -1,10 +1,34 @@
-from typing import Optional, Tuple
+"""Normalization and denormalization transforms for image patches."""
+
+from typing import Optional
 
 import numpy as np
+from numpy.typing import NDArray
 
 from careamics.transforms.transform import Transform
 
 
+def _reshape_stats(stats: list[float], ndim: int) -> NDArray:
+    """Reshape stats to match the number of dimensions of the input image.
+
+    This allows to broadcast the stats (mean or std) to the image dimensions, and
+    thus directly perform a vectorial calculation.
+
+    Parameters
+    ----------
+    stats : list of float
+        List of stats, mean or standard deviation.
+    ndim : int
+        Number of dimensions of the image, including the C channel.
+
+    Returns
+    -------
+    NDArray
+        Reshaped stats.
+    """
+    return np.array(stats)[(..., *[np.newaxis] * (ndim - 1))]
+
+
 class Normalize(Transform):
     """
     Normalize an image or image patch.
@@ -15,106 +39,205 @@ class Normalize(Transform):
     Not that an epsilon value of 1e-6 is added to the standard deviation to avoid
     division by zero and that it returns a float32 image.
 
+    Parameters
+    ----------
+    image_means : list of float
+        Mean value per channel.
+    image_stds : list of float
+        Standard deviation value per channel.
+    target_means : list of float, optional
+        Target mean value per channel, by default None.
+    target_stds : list of float, optional
+        Target standard deviation value per channel, by default None.
+
     Attributes
     ----------
-    mean : float
-        Mean value.
-    std : float
-        Standard deviation value.
+    image_means : list of float
+        Mean value per channel.
+    image_stds : list of float
+        Standard deviation value per channel.
+    target_means :list of float, optional
+        Target mean value per channel, by default None.
+    target_stds : list of float, optional
+        Target standard deviation value per channel, by default None.
     """
 
     def __init__(
         self,
-        mean: float,
-        std: float,
+        image_means: list[float],
+        image_stds: list[float],
+        target_means: Optional[list[float]] = None,
+        target_stds: Optional[list[float]] = None,
     ):
-        self.mean = mean
-        self.std = std
+        """Constructor.
+
+        Parameters
+        ----------
+        image_means : list of float
+            Mean value per channel.
+        image_stds : list of float
+            Standard deviation value per channel.
+        target_means : list of float, optional
+            Target mean value per channel, by default None.
+        target_stds : list of float, optional
+            Target standard deviation value per channel, by default None.
+        """
+        self.image_means = image_means
+        self.image_stds = image_stds
+        self.target_means = target_means
+        self.target_stds = target_stds
+
         self.eps = 1e-6
 
     def __call__(
-        self, patch: np.ndarray, target: Optional[np.ndarray] = None
-    ) -> Tuple[np.ndarray, Optional[np.ndarray]]:
+        self, patch: np.ndarray, target: Optional[NDArray] = None
+    ) -> tuple[NDArray, Optional[NDArray]]:
         """Apply the transform to the source patch and the target (optional).
 
         Parameters
         ----------
-        patch : np.ndarray
+        patch : NDArray
             Patch, 2D or 3D, shape C(Z)YX.
-        target : Optional[np.ndarray], optional
-            Target for the patch, by default None
+        target : NDArray, optional
+            Target for the patch, by default None.
 
         Returns
         -------
-        Tuple[np.ndarray, Optional[np.ndarray]]
-            Transformed patch and target.
+        tuple of NDArray
+            Transformed patch and target, the target can be returned as `None`.
         """
-        norm_patch = self._apply(patch)
-        norm_target = self._apply(target) if target is not None else None
+        if len(self.image_means) != patch.shape[0]:
+            raise ValueError(
+                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."
+            )
+
+        # reshape mean and std and apply the normalization to the patch
+        means = _reshape_stats(self.image_means, patch.ndim)
+        stds = _reshape_stats(self.image_stds, patch.ndim)
+        norm_patch = self._apply(patch, means, stds)
+
+        # same for the target patch
+        if (
+            target is not None
+            and self.target_means is not None
+            and self.target_stds is not None
+        ):
+            target_means = _reshape_stats(self.target_means, target.ndim)
+            target_stds = _reshape_stats(self.target_stds, target.ndim)
+            norm_target = self._apply(target, target_means, target_stds)
+        else:
+            norm_target = None
 
         return norm_patch, norm_target
 
-    def _apply(self, patch: np.ndarray) -> np.ndarray:
-        return ((patch - self.mean) / (self.std + self.eps)).astype(np.float32)
+    def _apply(self, patch: NDArray, mean: NDArray, std: NDArray) -> NDArray:
+        """
+        Apply the transform to the image.
+
+        Parameters
+        ----------
+        patch : NDArray
+            Image patch, 2D or 3D, shape C(Z)YX.
+        mean : NDArray
+            Mean values.
+        std : NDArray
+            Standard deviations.
+
+        Returns
+        -------
+        NDArray
+            Normalized image patch.
+        """
+        return ((patch - mean) / (std + self.eps)).astype(np.float32)
 
 
 class Denormalize:
     """
-    Denormalize an image or image patch.
+    Denormalize an image.
 
     Denormalization is performed expecting a zero mean and unit variance input. This
     transform expects C(Z)YX dimensions.
 
-    Not that an epsilon value of 1e-6 is added to the standard deviation to avoid
+    Note that an epsilon value of 1e-6 is added to the standard deviation to avoid
     division by zero during the normalization step, which is taken into account during
     denormalization.
 
-    Attributes
+    Parameters
     ----------
-    mean : float
-        Mean value.
-    std : float
-        Standard deviation value.
+    image_means : list or tuple of float
+        Mean value per channel.
+    image_stds : list or tuple of float
+        Standard deviation value per channel.
+
     """
 
     def __init__(
         self,
-        mean: float,
-        std: float,
+        image_means: list[float],
+        image_stds: list[float],
     ):
-        self.mean = mean
-        self.std = std
+        """Constructor.
+
+        Parameters
+        ----------
+        image_means : list of float
+            Mean value per channel.
+        image_stds : list of float
+            Standard deviation value per channel.
+        """
+        self.image_means = image_means
+        self.image_stds = image_stds
+
         self.eps = 1e-6
 
-    def __call__(
-        self, patch: np.ndarray, target: Optional[np.ndarray] = None
-    ) -> Tuple[np.ndarray, Optional[np.ndarray]]:
-        """Apply the transform to the source patch and the target (optional).
+    def __call__(self, patch: NDArray) -> NDArray:
+        """Reverse the normalization operation for a batch of patches.
 
         Parameters
         ----------
-        patch : np.ndarray
-            Patch, 2D or 3D, shape C(Z)YX.
-        target : Optional[np.ndarray], optional
-            Target for the patch, by default None
+        patch : NDArray
+            Patch, 2D or 3D, shape BC(Z)YX.
 
         Returns
         -------
-        Tuple[np.ndarray, Optional[np.ndarray]]
-            Transformed patch and target.
+        NDArray
+            Transformed array.
         """
-        norm_patch = self._apply(patch)
-        norm_target = self._apply(target) if target is not None else None
+        if len(self.image_means) != patch.shape[1]:
+            raise ValueError(
+                f"Number of means (got a list of size {len(self.image_means)}) and "
+                f"number of channels (got shape {patch.shape} for BC(Z)YX) do not "
+                f"match."
+            )
 
-        return norm_patch, norm_target
+        means = _reshape_stats(self.image_means, patch.ndim)
+        stds = _reshape_stats(self.image_stds, patch.ndim)
 
-    def _apply(self, patch: np.ndarray) -> np.ndarray:
+        denorm_array = self._apply(
+            patch,
+            np.swapaxes(means, 0, 1),  # swap axes as C channel is axis 1
+            np.swapaxes(stds, 0, 1),
+        )
+
+        return denorm_array.astype(np.float32)
+
+    def _apply(self, array: NDArray, mean: NDArray, std: NDArray) -> NDArray:
         """
         Apply the transform to the image.
 
         Parameters
         ----------
-        patch : np.ndarray
-            Image or image patch, 2D or 3D, shape C(Z)YX.
+        array : NDArray
+            Image patch, 2D or 3D, shape C(Z)YX.
+        mean : NDArray
+            Mean values.
+        std : NDArray
+            Standard deviations.
+
+        Returns
+        -------
+        NDArray
+            Denormalized image array.
         """
-        return patch * (self.std + self.eps) + self.mean
+        return array * (std + self.eps) + mean

careamics/transforms/pixel_manipulation.py

@@ -5,7 +5,7 @@ Pixel manipulation is used in N2V and similar algorithm to replace the value of
 masked pixels.
 """
 
-from typing import Optional, Tuple, Union
+from typing import Optional, Tuple
 
 import numpy as np
 
@@ -13,9 +13,12 @@ from .struct_mask_parameters import StructMaskParameters
 
 
 def _apply_struct_mask(
-    patch: np.ndarray, coords: np.ndarray, struct_params: StructMaskParameters
+    patch: np.ndarray,
+    coords: np.ndarray,
+    struct_params: StructMaskParameters,
+    rng: Optional[np.random.Generator] = None,
 ) -> np.ndarray:
-    """Applies structN2V masks to patch.
+    """Apply structN2V masks to patch.
 
     Each point in `coords` corresponds to the center of a mask, masks are paremeterized
     by `struct_params` and pixels in the mask (with respect to `coords`) are replaced by
@@ -31,12 +34,17 @@ def _apply_struct_mask(
         Coordinates of the ROI(subpatch) centers.
     struct_params : StructMaskParameters
         Parameters for the structN2V mask (axis and span).
+    rng : np.random.Generator or None
+        Random number generator.
 
     Returns
     -------
     np.ndarray
         Patch with the structN2V mask applied.
     """
+    if rng is None:
+        rng = np.random.default_rng()
+
     # relative axis
     moving_axis = -1 - struct_params.axis
 
@@ -67,7 +75,7 @@ def _apply_struct_mask(
     mix = np.delete(mix, mix[:, moving_axis] > max_bound, axis=0)
 
     # replace neighbouring pixels with random values from flat dist
-    patch[tuple(mix.T)] = np.random.uniform(patch.min(), patch.max(), size=mix.shape[0])
+    patch[tuple(mix.T)] = rng.uniform(patch.min(), patch.max(), size=mix.shape[0])
 
     return patch
 
@@ -98,7 +106,9 @@ def _odd_jitter_func(step: float, rng: np.random.Generator) -> np.ndarray:
 
 
 def _get_stratified_coords(
-    mask_pixel_perc: float, shape: Union[Tuple[int, int], Tuple[int, int, int]]
+    mask_pixel_perc: float,
+    shape: Tuple[int, ...],
+    rng: Optional[np.random.Generator] = None,
 ) -> np.ndarray:
     """
     Generate coordinates of the pixels to mask.
@@ -113,6 +123,8 @@ def _get_stratified_coords(
         calculating the distance between masked pixels across each axis.
     shape : Tuple[int, ...]
         Shape of the input patch.
+    rng : np.random.Generator or None
+        Random number generator.
 
     Returns
     -------
@@ -124,7 +136,8 @@ def _get_stratified_coords(
             "Calculating coordinates is only possible for 2D and 3D patches"
         )
 
-    rng = np.random.default_rng()
+    if rng is None:
+        rng = np.random.default_rng()
 
     mask_pixel_distance = np.round((100 / mask_pixel_perc) ** (1 / len(shape))).astype(
         np.int32
@@ -228,6 +241,7 @@ def uniform_manipulate(
     subpatch_size: int = 11,
     remove_center: bool = True,
     struct_params: Optional[StructMaskParameters] = None,
+    rng: Optional[np.random.Generator] = None,
 ) -> Tuple[np.ndarray, np.ndarray]:
     """
     Manipulate pixels by replacing them with a neighbor values.
@@ -248,19 +262,23 @@ def uniform_manipulate(
         Size of the subpatch the new pixel value is sampled from, by default 11.
     remove_center : bool
         Whether to remove the center pixel from the subpatch, by default False.
-    struct_params: Optional[StructMaskParameters]
+    struct_params : StructMaskParameters or None
         Parameters for the structN2V mask (axis and span).
+    rng : np.random.Generator or None
+        Random number generator.
 
     Returns
     -------
     Tuple[np.ndarray]
         Tuple containing the manipulated patch and the corresponding mask.
     """
+    if rng is None:
+        rng = np.random.default_rng()
+
     # Get the coordinates of the pixels to be replaced
     transformed_patch = patch.copy()
 
-    subpatch_centers = _get_stratified_coords(mask_pixel_percentage, patch.shape)
-    rng = np.random.default_rng()
+    subpatch_centers = _get_stratified_coords(mask_pixel_percentage, patch.shape, rng)
 
     # Generate coordinate grid for subpatch
     roi_span_full = np.arange(
@@ -303,6 +321,7 @@ def median_manipulate(
     mask_pixel_percentage: float,
     subpatch_size: int = 11,
     struct_params: Optional[StructMaskParameters] = None,
+    rng: Optional[np.random.Generator] = None,
 ) -> Tuple[np.ndarray, np.ndarray]:
     """
     Manipulate pixels by replacing them with the median of their surrounding subpatch.
@@ -322,18 +341,23 @@ def median_manipulate(
         Approximate percentage of pixels to be masked.
     subpatch_size : int
         Size of the subpatch the new pixel value is sampled from, by default 11.
-    struct_params: Optional[StructMaskParameters]
+    struct_params : StructMaskParameters or None, optional
         Parameters for the structN2V mask (axis and span).
+    rng : np.random.Generator or None, optional
+        Random number generato, by default None.
 
     Returns
     -------
     Tuple[np.ndarray]
            Tuple containing the manipulated patch, the original patch and the mask.
     """
+    if rng is None:
+        rng = np.random.default_rng()
+
     transformed_patch = patch.copy()
 
     # Get the coordinates of the pixels to be replaced
-    subpatch_centers = _get_stratified_coords(mask_pixel_percentage, patch.shape)
+    subpatch_centers = _get_stratified_coords(mask_pixel_percentage, patch.shape, rng)
 
     # Generate coordinate grid for subpatch
     roi_span = np.array(

careamics/transforms/struct_mask_parameters.py

@@ -1,3 +1,5 @@
+"""Class representing the parameters of structN2V masks."""
+
 from dataclasses import dataclass
 from typing import Literal
 
@@ -6,7 +8,7 @@ from typing import Literal
 class StructMaskParameters:
     """Parameters of structN2V masks.
 
-    Parameters
+    Attributes
     ----------
     axis : Literal[0, 1]
         Axis along which to apply the mask, horizontal (0) or vertical (1).

careamics/transforms/transform.py

@@ -1,33 +1,24 @@
 """A general parent class for transforms."""
 
-from typing import Optional, Tuple
-
-import numpy as np
+from typing import Any
 
 
 class Transform:
     """A general parent class for transforms."""
 
-    def __call__(
-        self, patch: np.ndarray, target: Optional[np.ndarray] = None
-    ) -> Tuple[np.ndarray, ...]:
-        """Apply the transform to the input data.
+    def __call__(self, *args: Any, **kwargs: Any) -> Any:
+        """Apply the transform.
 
         Parameters
         ----------
-        patch : np.ndarray
-            The input data to transform.
-        target : Optional[np.ndarray], optional
-            The target data to transform, by default None
+        *args : Any
+            Arguments.
+        **kwargs : Any
+            Keyword arguments.
 
         Returns
         -------
-        Tuple[np.ndarray, ...]
-            The output of the transformations.
-
-        Raises
-        ------
-        NotImplementedError
-            This method should be implemented in the child class.
+        Any
+            Transformed data.
         """
-        raise NotImplementedError
+        pass

careamics/transforms/tta.py

@@ -1,11 +1,8 @@
 """Test-time augmentations."""
 
-from typing import List
-
 from torch import Tensor, flip, mean, rot90, stack
 
 
-# TODO add tests
 class ImageRestorationTTA:
     """
     Test-time augmentation for image restoration tasks.
@@ -13,62 +10,79 @@ class ImageRestorationTTA:
     The augmentation is performed using all 90 deg rotations and their flipped version,
     as well as the original image flipped.
 
-    Tensors should be of shape SC(Z)YX
+    Tensors should be of shape SC(Z)YX.
 
     This transformation is used in the LightningModule in order to perform test-time
-    agumentation.
+    augmentation.
     """
 
-    def __init__(self) -> None:
-        """Constructor."""
-        pass
-
-    def forward(self, x: Tensor) -> List[Tensor]:
+    def forward(self, input_tensor: Tensor) -> list[Tensor]:
         """
         Apply test-time augmentation to the input tensor.
 
         Parameters
         ----------
-        x : Tensor
+        input_tensor : Tensor
             Input tensor, shape SC(Z)YX.
 
         Returns
         -------
-        List[Tensor]
+        list of torch.Tensor
             List of augmented tensors.
         """
+        # axes: only applies to YX axes
+        axes = (-2, -1)
+
         augmented = [
-            x,
-            rot90(x, 1, dims=(-2, -1)),
-            rot90(x, 2, dims=(-2, -1)),
-            rot90(x, 3, dims=(-2, -1)),
+            # original
+            input_tensor,
+            # rotations
+            rot90(input_tensor, 1, dims=axes),
+            rot90(input_tensor, 2, dims=axes),
+            rot90(input_tensor, 3, dims=axes),
+            # original flipped
+            flip(input_tensor, dims=(axes[0],)),
+            flip(input_tensor, dims=(axes[1],)),
         ]
-        augmented_flip = augmented.copy()
-        for x_ in augmented:
-            augmented_flip.append(flip(x_, dims=(-3, -1)))
-        return augmented_flip
 
-    def backward(self, x: List[Tensor]) -> Tensor:
+        # rotated once, flipped
+        augmented.extend(
+            [
+                flip(augmented[1], dims=(axes[0],)),
+                flip(augmented[1], dims=(axes[1],)),
+            ]
+        )
+
+        return augmented
+
+    def backward(self, x: list[Tensor]) -> Tensor:
         """Undo the test-time augmentation.
 
         Parameters
         ----------
         x : Any
-            List of augmented tensors.
+            List of augmented tensors of shape SC(Z)YX.
 
         Returns
         -------
         Any
             Original tensor.
         """
+        axes = (-2, -1)
+
         reverse = [
+            # original
             x[0],
-            rot90(x[1], -1, dims=(-2, -1)),
-            rot90(x[2], -2, dims=(-2, -1)),
-            rot90(x[3], -3, dims=(-2, -1)),
-            flip(x[4], dims=(-3, -1)),
-            rot90(flip(x[5], dims=(-3, -1)), -1, dims=(-2, -1)),
-            rot90(flip(x[6], dims=(-3, -1)), -2, dims=(-2, -1)),
-            rot90(flip(x[7], dims=(-3, -1)), -3, dims=(-2, -1)),
+            # rotated
+            rot90(x[1], -1, dims=axes),
+            rot90(x[2], -2, dims=axes),
+            rot90(x[3], -3, dims=axes),
+            # original flipped
+            flip(x[4], dims=(axes[0],)),
+            flip(x[5], dims=(axes[1],)),
+            # rotated once, flipped
+            rot90(flip(x[6], dims=(axes[0],)), -1, dims=axes),
+            rot90(flip(x[7], dims=(axes[1],)), -1, dims=axes),
         ]
+
         return mean(stack(reverse), dim=0)