careamics 0.1.0rc4__py3-none-any.whl → 0.1.0rc6__py3-none-any.whl

careamics/transforms/compose.py

@@ -0,0 +1,124 @@
+"""A class chaining transforms together."""
+
+from typing import Callable, Dict, List, Optional, Tuple
+
+import numpy as np
+
+from careamics.config.data_model import TRANSFORMS_UNION
+
+from .n2v_manipulate import N2VManipulate
+from .normalize import Normalize
+from .transform import Transform
+from .xy_flip import XYFlip
+from .xy_random_rotate90 import XYRandomRotate90
+
+ALL_TRANSFORMS = {
+    "Normalize": Normalize,
+    "N2VManipulate": N2VManipulate,
+    "XYFlip": XYFlip,
+    "XYRandomRotate90": XYRandomRotate90,
+}
+
+
+def get_all_transforms() -> Dict[str, type]:
+    """Return all the transforms accepted by CAREamics.
+
+    Returns
+    -------
+    dict
+        A dictionary with all the transforms accepted by CAREamics, where the keys are
+        the transform names and the values are the transform classes.
+    """
+    return ALL_TRANSFORMS
+
+
+class Compose:
+    """A class chaining transforms together.
+
+    Parameters
+    ----------
+    transform_list : List[TRANSFORMS_UNION]
+        A list of dictionaries where each dictionary contains the name of a
+        transform and its parameters.
+
+    Attributes
+    ----------
+    _callable_transforms : Callable
+        A callable that applies the transforms to the input data.
+    """
+
+    def __init__(self, transform_list: List[TRANSFORMS_UNION]) -> None:
+        """Instantiate a Compose object.
+
+        Parameters
+        ----------
+        transform_list : List[TRANSFORMS_UNION]
+            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()
+
+        # instantiate all transforms
+        transforms = [all_transforms[t.name](**t.model_dump()) for t in transform_list]
+
+        self._callable_transforms = self._chain_transforms(transforms)
+
+    def _chain_transforms(self, transforms: List[Transform]) -> Callable:
+        """Chain the transforms together.
+
+        Parameters
+        ----------
+        transforms : List[Transform]
+            A list of transforms to chain together.
+
+        Returns
+        -------
+        Callable
+            A callable that applies the transforms in order to the input data.
+        """
+
+        def _chain(
+            patch: np.ndarray, target: Optional[np.ndarray]
+        ) -> Tuple[np.ndarray, Optional[np.ndarray]]:
+            """Chain transforms on the input data.
+
+            Parameters
+            ----------
+            patch : np.ndarray
+                Input data.
+            target : Optional[np.ndarray]
+                Target data, by default None.
+
+            Returns
+            -------
+            Tuple[np.ndarray, Optional[np.ndarray]]
+                The output of the transformations.
+            """
+            params = (patch, target)
+
+            for t in transforms:
+                params = t(*params)
+
+            return params
+
+        return _chain
+
+    def __call__(
+        self, patch: np.ndarray, target: Optional[np.ndarray] = None
+    ) -> Tuple[np.ndarray, ...]:
+        """Apply the transforms to the input data.
+
+        Parameters
+        ----------
+        patch : np.ndarray
+            The input data.
+        target : Optional[np.ndarray], optional
+            Target data, by default None.
+
+        Returns
+        -------
+        Tuple[np.ndarray, ...]
+            The output of the transformations.
+        """
+        return self._callable_transforms(patch, target)

careamics/transforms/n2v_manipulate.py

@@ -1,26 +1,53 @@
+"""N2V manipulation transform."""
+
 from typing import Any, Literal, Optional, Tuple
 
 import numpy as np
-from albumentations import ImageOnlyTransform
 
 from careamics.config.support import SupportedPixelManipulation, SupportedStructAxis
+from careamics.transforms.transform import Transform
 
 from .pixel_manipulation import median_manipulate, uniform_manipulate
 from .struct_mask_parameters import StructMaskParameters
 
 
-class N2VManipulate(ImageOnlyTransform):
+class N2VManipulate(Transform):
     """
     Default augmentation for the N2V model.
 
-    This transform expects (Z)YXC dimensions.
+    This transform expects C(Z)YX dimensions.
 
     Parameters
     ----------
-    mask_pixel_percentage : float
-        Approximate percentage of pixels to be masked.
+    roi_size : int, optional
+        Size of the replacement area, by default 11.
+    masked_pixel_percentage : float, optional
+        Percentage of pixels to mask, by default 0.2.
+    strategy : Literal[ "uniform", "median" ], optional
+        Replaccement strategy, uniform or median, by default uniform.
+    remove_center : bool, optional
+        Whether to remove central pixel from patch, by default True.
+    struct_mask_axis : Literal["horizontal", "vertical", "none"], optional
+        StructN2V mask axis, by default "none".
+    struct_mask_span : int, optional
+        StructN2V mask span, by default 5.
+    seed : Optional[int], optional
+        Random seed, by default None.
+
+    Attributes
+    ----------
+    masked_pixel_percentage : float
+        Percentage of pixels to mask.
     roi_size : int
-        Size of the ROI the new pixel value is sampled from, by default 11.
+        Size of the replacement area.
+    strategy : Literal[ "uniform", "median" ]
+        Replaccement strategy, uniform or median.
+    remove_center : bool
+        Whether to remove central pixel from patch.
+    struct_mask : Optional[StructMaskParameters]
+        StructN2V mask parameters.
+    rng : Generator
+        Random number generator.
     """
 
     def __init__(
@@ -33,29 +60,31 @@ class N2VManipulate(ImageOnlyTransform):
         remove_center: bool = True,
         struct_mask_axis: Literal["horizontal", "vertical", "none"] = "none",
         struct_mask_span: int = 5,
+        seed: Optional[int] = None,  # TODO use in pixel manipulation
     ):
         """Constructor.
 
         Parameters
         ----------
         roi_size : int, optional
-            Size of the replacement area, by default 11
+            Size of the replacement area, by default 11.
         masked_pixel_percentage : float, optional
-            Percentage of pixels to mask, by default 0.2
+            Percentage of pixels to mask, by default 0.2.
         strategy : Literal[ "uniform", "median" ], optional
-            Replaccement strategy, uniform or median, by default uniform
+            Replaccement strategy, uniform or median, by default uniform.
         remove_center : bool, optional
-            Whether to remove central pixel from patch, by default True
+            Whether to remove central pixel from patch, by default True.
         struct_mask_axis : Literal["horizontal", "vertical", "none"], optional
-            StructN2V mask axis, by default "none"
+            StructN2V mask axis, by default "none".
         struct_mask_span : int, optional
-            StructN2V mask span, by default 5
+            StructN2V mask span, by default 5.
+        seed : Optional[int], optional
+            Random seed, by default None.
         """
-        super().__init__(p=1)
         self.masked_pixel_percentage = masked_pixel_percentage
         self.roi_size = roi_size
         self.strategy = strategy
-        self.remove_center = remove_center
+        self.remove_center = remove_center  # TODO is this ever used?
 
         if struct_mask_axis == SupportedStructAxis.NONE:
             self.struct_mask: Optional[StructMaskParameters] = None
@@ -65,23 +94,35 @@ class N2VManipulate(ImageOnlyTransform):
                 span=struct_mask_span,
             )
 
-    def apply(
-        self, patch: np.ndarray, **kwargs: Any
+        # numpy random generator
+        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]:
         """Apply the transform to the image.
 
         Parameters
         ----------
-        image : np.ndarray
-            Image or image patch, 2D or 3D, shape (y, x, c) or (z, y, x, c).
+        patch : np.ndarray
+            Image patch, 2D or 3D, shape C(Z)YX.
+        *args : Any
+            Additional arguments, unused.
+        **kwargs : Any
+            Additional keyword arguments, unused.
+
+        Returns
+        -------
+        Tuple[np.ndarray, np.ndarray, np.ndarray]
+            Masked patch, original patch, and mask.
         """
         masked = np.zeros_like(patch)
         mask = np.zeros_like(patch)
         if self.strategy == SupportedPixelManipulation.UNIFORM:
             # Iterate over the channels to apply manipulation separately
-            for c in range(patch.shape[-1]):
-                masked[..., c], mask[..., c] = uniform_manipulate(
-                    patch=patch[..., c],
+            for c in range(patch.shape[0]):
+                masked[c, ...], mask[c, ...] = uniform_manipulate(
+                    patch=patch[c, ...],
                     mask_pixel_percentage=self.masked_pixel_percentage,
                     subpatch_size=self.roi_size,
                     remove_center=self.remove_center,
@@ -89,9 +130,9 @@ class N2VManipulate(ImageOnlyTransform):
                 )
         elif self.strategy == SupportedPixelManipulation.MEDIAN:
             # Iterate over the channels to apply manipulation separately
-            for c in range(patch.shape[-1]):
-                masked[..., c], mask[..., c] = median_manipulate(
-                    patch=patch[..., c],
+            for c in range(patch.shape[0]):
+                masked[c, ...], mask[c, ...] = median_manipulate(
+                    patch=patch[c, ...],
                     mask_pixel_percentage=self.masked_pixel_percentage,
                     subpatch_size=self.roi_size,
                     struct_params=self.struct_mask,
@@ -101,13 +142,3 @@ class N2VManipulate(ImageOnlyTransform):
 
         # TODO why return patch?
         return masked, patch, mask
-
-    def get_transform_init_args_names(self) -> Tuple[str, ...]:
-        """Get the transform parameters.
-
-        Returns
-        -------
-        Tuple[str, ...]
-            Transform parameters.
-        """
-        return ("roi_size", "masked_pixel_percentage", "strategy", "struct_mask")

careamics/transforms/normalize.py

@@ -1,27 +1,35 @@
-from typing import Any
+"""Normalization and denormalization transforms for image patches."""
+
+from typing import Optional, Tuple
 
 import numpy as np
-from albumentations import DualTransform
+
+from careamics.transforms.transform import Transform
 
 
-class Normalize(DualTransform):
+class Normalize(Transform):
     """
     Normalize an image or image patch.
 
-    Normalization is a zero mean and unit variance. This transform expects (Z)YXC
+    Normalization is a zero mean and unit variance. This transform expects C(Z)YX
     dimensions.
 
     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
+    ----------
+    mean : float
+        Mean value.
+    std : float
+        Standard deviation value.
+
     Attributes
     ----------
     mean : float
         Mean value.
     std : float
         Standard deviation value.
-    eps : float
-        Epsilon value to avoid division by zero.
     """
 
     def __init__(
@@ -29,61 +37,82 @@ class Normalize(DualTransform):
         mean: float,
         std: float,
     ):
-        super().__init__(always_apply=True, p=1)
+        """Constructor.
 
+        Parameters
+        ----------
+        mean : float
+            Mean value.
+        std : float
+            Standard deviation value.
+        """
         self.mean = mean
         self.std = std
         self.eps = 1e-6
 
-    def apply(self, patch: np.ndarray, **kwargs: Any) -> np.ndarray:
-        """
-        Apply the transform to the image.
+    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).
 
         Parameters
         ----------
         patch : np.ndarray
-            Image or image patch, 2D or 3D, shape (y, x, c) or (z, y, x, c).
+            Patch, 2D or 3D, shape C(Z)YX.
+        target : Optional[np.ndarray], optional
+            Target for the patch, by default None.
 
         Returns
         -------
-        np.ndarray
-            Normalized image or image patch.
+        Tuple[np.ndarray, Optional[np.ndarray]]
+            Transformed patch and target.
         """
-        return ((patch - self.mean) / (self.std + self.eps)).astype(np.float32)
+        norm_patch = self._apply(patch)
+        norm_target = self._apply(target) if target is not None else None
 
-    def apply_to_mask(self, mask: np.ndarray, **kwargs: Any) -> np.ndarray:
-        """
-        Apply the transform to the mask.
+        return norm_patch, norm_target
 
-        The mask is returned as is.
+    def _apply(self, patch: np.ndarray) -> np.ndarray:
+        """
+        Apply the transform to the image.
 
         Parameters
         ----------
-        mask : np.ndarray
-            Mask or mask patch, 2D or 3D, shape (y, x, c) or (z, y, x, c).
+        patch : np.ndarray
+            Image patch, 2D or 3D, shape C(Z)YX.
+
+        Returns
+        -------
+        np.ndarray
+            Normalizedimage patch.
         """
-        return mask
+        return ((patch - self.mean) / (self.std + self.eps)).astype(np.float32)
 
 
-class Denormalize(DualTransform):
+class Denormalize:
     """
     Denormalize an image or image patch.
 
     Denormalization is performed expecting a zero mean and unit variance input. This
-    transform expects (Z)YXC dimensions.
+    transform expects C(Z)YX dimensions.
 
     Not 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.
 
+    Parameters
+    ----------
+    mean : float
+        Mean value.
+    std : float
+        Standard deviation value.
+
     Attributes
     ----------
     mean : float
         Mean value.
     std : float
         Standard deviation value.
-    eps : float
-        Epsilon value to avoid division by zero.
     """
 
     def __init__(
@@ -91,19 +120,53 @@ class Denormalize(DualTransform):
         mean: float,
         std: float,
     ):
-        super().__init__(always_apply=True, p=1)
+        """Constructor.
 
+        Parameters
+        ----------
+        mean : float
+            Mean.
+        std : float
+            Standard deviation.
+        """
         self.mean = mean
         self.std = std
         self.eps = 1e-6
 
-    def apply(self, patch: np.ndarray, **kwargs: Any) -> np.ndarray:
+    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).
+
+        Parameters
+        ----------
+        patch : np.ndarray
+            Patch, 2D or 3D, shape C(Z)YX.
+        target : Optional[np.ndarray], optional
+            Target for the patch, by default None.
+
+        Returns
+        -------
+        Tuple[np.ndarray, Optional[np.ndarray]]
+            Transformed patch and target.
+        """
+        norm_patch = self._apply(patch)
+        norm_target = self._apply(target) if target is not None else None
+
+        return norm_patch, norm_target
+
+    def _apply(self, patch: np.ndarray) -> np.ndarray:
         """
         Apply the transform to the image.
 
         Parameters
         ----------
         patch : np.ndarray
-            Image or image patch, 2D or 3D, shape (y, x, c) or (z, y, x, c).
+            Image patch, 2D or 3D, shape C(Z)YX.
+
+        Returns
+        -------
+        np.ndarray
+            Denormalized image patch.
         """
         return patch * (self.std + self.eps) + self.mean

careamics/transforms/pixel_manipulation.py

@@ -4,7 +4,8 @@ Pixel manipulation methods.
 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
 
@@ -14,7 +15,7 @@ from .struct_mask_parameters import StructMaskParameters
 def _apply_struct_mask(
     patch: np.ndarray, coords: np.ndarray, struct_params: StructMaskParameters
 ) -> 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
@@ -97,7 +98,7 @@ 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, ...]
 ) -> np.ndarray:
     """
     Generate coordinates of the pixels to mask.
@@ -246,9 +247,8 @@ def uniform_manipulate(
     subpatch_size : int
         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. See
-        uniform with/without central pixel in the documentation. #TODO add link
-    struct_params: Optional[StructMaskParameters]
+        Whether to remove the center pixel from the subpatch, by default False.
+    struct_params : Optional[StructMaskParameters]
         Parameters for the structN2V mask (axis and span).
 
     Returns
@@ -322,7 +322,7 @@ 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 : Optional[StructMaskParameters]
         Parameters for the structN2V mask (axis and span).
 
     Returns

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

@@ -0,0 +1,24 @@
+"""A general parent class for transforms."""
+
+from typing import Any
+
+
+class Transform:
+    """A general parent class for transforms."""
+
+    def __call__(self, *args: Any, **kwargs: Any) -> Any:
+        """Apply the transform.
+
+        Parameters
+        ----------
+        *args : Any
+            Arguments.
+        **kwargs : Any
+            Keyword arguments.
+
+        Returns
+        -------
+        Any
+            Transformed data.
+        """
+        pass

careamics/transforms/tta.py

@@ -1,7 +1,7 @@
 """Test-time augmentations."""
+
 from typing import List
 
-import numpy as np
 from torch import Tensor, flip, mean, rot90, stack
 
 
@@ -48,7 +48,7 @@ class ImageRestorationTTA:
             augmented_flip.append(flip(x_, dims=(-3, -1)))
         return augmented_flip
 
-    def backward(self, x: List[Tensor]) -> np.ndarray:
+    def backward(self, x: List[Tensor]) -> Tensor:
         """Undo the test-time augmentation.
 
         Parameters