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

careamics/models/unet.py

@@ -34,7 +34,9 @@ class UnetEncoder(nn.Module):
         Dropout probability, by default 0.0.
     pool_kernel : int, optional
         Kernel size for the max pooling layers, by default 2.
-    groups: int, optional
+    n2v2 : bool, optional
+        Whether to use N2V2 architecture, by default False.
+    groups : int, optional
         Number of blocked connections from input channels to output
         channels, by default 1.
     """
@@ -70,7 +72,9 @@ class UnetEncoder(nn.Module):
             Dropout probability, by default 0.0.
         pool_kernel : int, optional
             Kernel size for the max pooling layers, by default 2.
-        groups: int, optional
+        n2v2 : bool, optional
+            Whether to use N2V2 architecture, by default False.
+        groups : int, optional
             Number of blocked connections from input channels to output
             channels, by default 1.
         """
@@ -140,7 +144,9 @@ class UnetDecoder(nn.Module):
         Whether to use batch normalization, by default True.
     dropout : float, optional
         Dropout probability, by default 0.0.
-    groups: int, optional
+    n2v2 : bool, optional
+        Whether to use N2V2 architecture, by default False.
+    groups : int, optional
         Number of blocked connections from input channels to output
         channels, by default 1.
     """
@@ -170,7 +176,9 @@ class UnetDecoder(nn.Module):
             Whether to use batch normalization, by default True.
         dropout : float, optional
             Dropout probability, by default 0.0.
-        groups: int, optional
+        n2v2 : bool, optional
+            Whether to use N2V2 architecture, by default False.
+        groups : int, optional
             Number of blocked connections from input channels to output
             channels, by default 1.
         """
@@ -250,22 +258,25 @@ class UnetDecoder(nn.Module):
 
     @staticmethod
     def _interleave(A: torch.Tensor, B: torch.Tensor, groups: int) -> torch.Tensor:
-        """
-        Splits the tensors `A` and `B` into equally sized groups along the
-        channel axis (axis=1); then concatenates the groups in alternating
-        order along the channel axis, starting with the first group from tensor
-        A.
+        """Interleave two tensors.
+
+        Splits the tensors `A` and `B` into equally sized groups along the channel
+        axis (axis=1); then concatenates the groups in alternating order along the
+        channel axis, starting with the first group from tensor A.
 
         Parameters
         ----------
-        A: torch.Tensor
-        B: torch.Tensor
-        groups: int
+        A : torch.Tensor
+            First tensor.
+        B : torch.Tensor
+            Second tensor.
+        groups : int
             The number of groups.
 
         Returns
         -------
         torch.Tensor
+            Interleaved tensor.
 
         Raises
         ------
@@ -322,8 +333,14 @@ class UNet(nn.Module):
         Dropout probability, by default 0.0.
     pool_kernel : int, optional
         Kernel size of the pooling layers, by default 2.
-    last_activation : Optional[Callable], optional
+    final_activation : Optional[Callable], optional
         Activation function to use for the last layer, by default None.
+    n2v2 : bool, optional
+        Whether to use N2V2 architecture, by default False.
+    independent_channels : bool
+        Whether to train the channels independently, by default True.
+    **kwargs : Any
+        Additional keyword arguments, unused.
     """
 
     def __init__(
@@ -362,11 +379,15 @@ class UNet(nn.Module):
             Dropout probability, by default 0.0.
         pool_kernel : int, optional
             Kernel size of the pooling layers, by default 2.
-        last_activation : Optional[Callable], optional
+        final_activation : Optional[Callable], optional
             Activation function to use for the last layer, by default None.
+        n2v2 : bool, optional
+            Whether to use N2V2 architecture, by default False.
         independent_channels : bool
             Whether to train parallel independent networks for each channel, by
             default True.
+        **kwargs : Any
+            Additional keyword arguments, unused.
         """
         super().__init__()
 

careamics/prediction/stitch_prediction.py

@@ -1,8 +1,4 @@
-"""
-Prediction convenience functions.
-
-These functions are used during prediction.
-"""
+"""Prediction utility functions."""
 
 from typing import List
 
@@ -21,7 +17,7 @@ def stitch_prediction(
     ----------
     tiles : List[torch.Tensor]
         Cropped tiles and their respective stitching coordinates.
-    stitching_coords : List
+    stitching_data : List
         List of information and coordinates obtained from
         `dataset.tiled_patching.extract_tiles`.
 

careamics/transforms/__init__.py

@@ -3,7 +3,7 @@
 __all__ = [
     "get_all_transforms",
     "N2VManipulate",
-    "NDFlip",
+    "XYFlip",
     "XYRandomRotate90",
     "ImageRestorationTTA",
     "Denormalize",
@@ -14,7 +14,7 @@ __all__ = [
 
 from .compose import Compose, get_all_transforms
 from .n2v_manipulate import N2VManipulate
-from .nd_flip import NDFlip
 from .normalize import Denormalize, Normalize
 from .tta import ImageRestorationTTA
+from .xy_flip import XYFlip
 from .xy_random_rotate90 import XYRandomRotate90

careamics/transforms/compose.py

@@ -1,26 +1,26 @@
 """A class chaining transforms together."""
 
-from typing import Callable, List, Optional, Tuple
+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 .nd_flip import NDFlip
 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,
-    "NDFlip": NDFlip,
+    "XYFlip": XYFlip,
     "XYRandomRotate90": XYRandomRotate90,
 }
 
 
-def get_all_transforms() -> dict:
+def get_all_transforms() -> Dict[str, type]:
     """Return all the transforms accepted by CAREamics.
 
     Returns
@@ -33,7 +33,19 @@ def get_all_transforms() -> dict:
 
 
 class Compose:
-    """A class chaining transforms together."""
+    """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.
@@ -68,7 +80,21 @@ class Compose:
 
         def _chain(
             patch: np.ndarray, target: Optional[np.ndarray]
-        ) -> Tuple[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:
@@ -88,7 +114,7 @@ class Compose:
         patch : np.ndarray
             The input data.
         target : Optional[np.ndarray], optional
-            Target data, by default None
+            Target data, by default None.
 
         Returns
         -------

careamics/transforms/n2v_manipulate.py

@@ -1,3 +1,5 @@
+"""N2V manipulation transform."""
+
 from typing import Any, Literal, Optional, Tuple
 
 import numpy as np
@@ -17,10 +19,35 @@ class N2VManipulate(Transform):
 
     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__(
@@ -40,24 +67,24 @@ class N2VManipulate(Transform):
         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
+            Random seed, by default None.
         """
         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
@@ -77,8 +104,17 @@ class N2VManipulate(Transform):
 
         Parameters
         ----------
-        image : np.ndarray
-            Image or image patch, 2D or 3D, shape C(Z)YX.
+        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)

careamics/transforms/normalize.py

@@ -1,3 +1,5 @@
+"""Normalization and denormalization transforms for image patches."""
+
 from typing import Optional, Tuple
 
 import numpy as np
@@ -15,6 +17,13 @@ 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
+    ----------
+    mean : float
+        Mean value.
+    std : float
+        Standard deviation value.
+
     Attributes
     ----------
     mean : float
@@ -28,6 +37,15 @@ class Normalize(Transform):
         mean: float,
         std: float,
     ):
+        """Constructor.
+
+        Parameters
+        ----------
+        mean : float
+            Mean value.
+        std : float
+            Standard deviation value.
+        """
         self.mean = mean
         self.std = std
         self.eps = 1e-6
@@ -42,7 +60,7 @@ class Normalize(Transform):
         patch : np.ndarray
             Patch, 2D or 3D, shape C(Z)YX.
         target : Optional[np.ndarray], optional
-            Target for the patch, by default None
+            Target for the patch, by default None.
 
         Returns
         -------
@@ -55,6 +73,19 @@ class Normalize(Transform):
         return norm_patch, norm_target
 
     def _apply(self, patch: np.ndarray) -> np.ndarray:
+        """
+        Apply the transform to the image.
+
+        Parameters
+        ----------
+        patch : np.ndarray
+            Image patch, 2D or 3D, shape C(Z)YX.
+
+        Returns
+        -------
+        np.ndarray
+            Normalizedimage patch.
+        """
         return ((patch - self.mean) / (self.std + self.eps)).astype(np.float32)
 
 
@@ -69,6 +100,13 @@ class Denormalize:
     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
@@ -82,6 +120,15 @@ class Denormalize:
         mean: float,
         std: float,
     ):
+        """Constructor.
+
+        Parameters
+        ----------
+        mean : float
+            Mean.
+        std : float
+            Standard deviation.
+        """
         self.mean = mean
         self.std = std
         self.eps = 1e-6
@@ -96,7 +143,7 @@ class Denormalize:
         patch : np.ndarray
             Patch, 2D or 3D, shape C(Z)YX.
         target : Optional[np.ndarray], optional
-            Target for the patch, by default None
+            Target for the patch, by default None.
 
         Returns
         -------
@@ -115,6 +162,11 @@ class Denormalize:
         Parameters
         ----------
         patch : np.ndarray
-            Image or image patch, 2D or 3D, shape C(Z)YX.
+            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

@@ -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
 
@@ -15,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
@@ -98,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.
@@ -248,7 +248,7 @@ 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 : 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

@@ -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/xy_flip.py

@@ -0,0 +1,123 @@
+"""XY flip transform."""
+
+from typing import Optional, Tuple
+
+import numpy as np
+
+from careamics.transforms.transform import Transform
+
+
+class XYFlip(Transform):
+    """Flip image along X and Y axis, one at a time.
+
+    This transform randomly flips one of the last two axes.
+
+    This transform expects C(Z)YX dimensions.
+
+    Attributes
+    ----------
+    axis_indices : List[int]
+        Indices of the axes that can be flipped.
+    rng : np.random.Generator
+        Random number generator.
+    p : float
+        Probability of applying the transform.
+    seed : Optional[int]
+        Random seed.
+
+    Parameters
+    ----------
+    flip_x : bool, optional
+        Whether to flip along the X axis, by default True.
+    flip_y : bool, optional
+        Whether to flip along the Y axis, by default True.
+    p : float, optional
+        Probability of applying the transform, by default 0.5.
+    seed : Optional[int], optional
+        Random seed, by default None.
+    """
+
+    def __init__(
+        self,
+        flip_x: bool = True,
+        flip_y: bool = True,
+        p: float = 0.5,
+        seed: Optional[int] = None,
+    ) -> None:
+        """Constructor.
+
+        Parameters
+        ----------
+        flip_x : bool, optional
+            Whether to flip along the X axis, by default True.
+        flip_y : bool, optional
+            Whether to flip along the Y axis, by default True.
+        p : float
+            Probability of applying the transform, by default 0.5.
+        seed : Optional[int], optional
+            Random seed, by default None.
+        """
+        if p < 0 or p > 1:
+            raise ValueError("Probability must be in [0, 1].")
+
+        if not flip_x and not flip_y:
+            raise ValueError("At least one axis must be flippable.")
+
+        # probability to apply the transform
+        self.p = p
+
+        # "flippable" axes
+        self.axis_indices = []
+
+        if flip_y:
+            self.axis_indices.append(-2)
+        if flip_x:
+            self.axis_indices.append(-1)
+
+        # numpy random generator
+        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]]:
+        """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.
+        """
+        if self.rng.random() > self.p:
+            return patch, target
+
+        # 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
+
+        return patch_transformed, target_transformed
+
+    def _apply(self, patch: np.ndarray, axis: int) -> np.ndarray:
+        """Apply the transform to the image.
+
+        Parameters
+        ----------
+        patch : np.ndarray
+            Image patch, 2D or 3D, shape C(Z)YX.
+        axis : int
+            Axis to flip.
+
+        Returns
+        -------
+        np.ndarray
+            Flipped image patch.
+        """
+        return np.ascontiguousarray(np.flip(patch, axis=axis))