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

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))

careamics/transforms/xy_random_rotate90.py

@@ -1,95 +1,101 @@
-from typing import Any, Dict, Tuple
+"""Patch transform applying XY random 90 degrees rotations."""
+
+from typing import Optional, Tuple
 
 import numpy as np
-from albumentations import DualTransform
+
+from careamics.transforms.transform import Transform
 
 
-class XYRandomRotate90(DualTransform):
+class XYRandomRotate90(Transform):
     """Applies random 90 degree rotations to the YX axis.
 
-    This transform expects (Z)YXC dimensions.
+    This transform expects C(Z)YX dimensions.
+
+    Attributes
+    ----------
+    rng : np.random.Generator
+        Random number generator.
+    p : float
+        Probability of applying the transform.
+    seed : Optional[int]
+        Random seed.
 
     Parameters
     ----------
-    p : int, optional
-        Probability to apply the transform, by default 0.5
-    is_3D : bool, optional
-        Whether the patches are 3D, by default False
+    p : float
+        Probability of applying the transform, by default 0.5.
+    seed : Optional[int]
+        Random seed, by default None.
     """
 
-    def __init__(self, p: float = 0.5, is_3D: bool = False):
+    def __init__(self, p: float = 0.5, seed: Optional[int] = None):
         """Constructor.
 
         Parameters
         ----------
-        p : float, optional
-            Probability to apply the transform, by default 0.5
-        is_3D : bool, optional
-            Whether the patches are 3D, by default False
+        p : float
+            Probability of applying the transform, by default 0.5.
+        seed : Optional[int]
+            Random seed, by default None.
         """
-        super().__init__(p=p)
+        if p < 0 or p > 1:
+            raise ValueError("Probability must be in [0, 1].")
+
+        # probability to apply the transform
+        self.p = p
 
-        self.is_3D = is_3D
+        # numpy random generator
+        self.rng = np.random.default_rng(seed=seed)
 
-        # rotation axes
-        if is_3D:
-            self.axes = (1, 2)
-        else:
-            self.axes = (0, 1)
+    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 get_params(self, **kwargs: Any) -> Dict[str, int]:
-        """Get the transform parameters.
+        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
         -------
-        Dict[str, int]
-            Transform parameters.
+        Tuple[np.ndarray, Optional[np.ndarray]]
+            Transformed patch and target.
         """
-        return {"n_rotations": np.random.randint(1, 4)}
+        if self.rng.random() > self.p:
+            return patch, target
 
-    def apply(self, patch: np.ndarray, n_rotations: int, **kwargs: Any) -> np.ndarray:
-        """Apply the transform to the image.
+        # number of rotations
+        n_rot = self.rng.integers(1, 4)
 
-        Parameters
-        ----------
-        patch : np.ndarray
-            Image or image patch, 2D or 3D, shape (y, x, c) or (z, y, x, c).
-        flip_axis : int
-            Axis along which to flip the patch.
-        """
-        if len(patch.shape) == 3 and self.is_3D:
-            raise ValueError(
-                "Incompatible patch shape and dimensionality. ZYXC patch shape "
-                "expected, but got YXC shape."
-            )
+        axes = (-2, -1)
+        patch_transformed = self._apply(patch, n_rot, axes)
+        target_transformed = (
+            self._apply(target, n_rot, axes) if target is not None else None
+        )
 
-        return np.ascontiguousarray(np.rot90(patch, k=n_rotations, axes=self.axes))
+        return patch_transformed, target_transformed
 
-    def apply_to_mask(
-        self, mask: np.ndarray, n_rotations: int, **kwargs: Any
+    def _apply(
+        self, patch: np.ndarray, n_rot: int, axes: Tuple[int, int]
     ) -> np.ndarray:
-        """Apply the transform to the mask.
+        """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).
-        """
-        if len(mask.shape) != 4 and self.is_3D:
-            raise ValueError(
-                "Incompatible mask shape and dimensionality. ZYXC patch shape "
-                "expected, but got YXC shape."
-            )
-
-        return np.ascontiguousarray(np.rot90(mask, k=n_rotations, axes=self.axes))
-
-    def get_transform_init_args_names(self) -> Tuple[str, str]:
-        """
-        Get the transform arguments.
+        patch : np.ndarray
+            Image or image patch, 2D or 3D, shape C(Z)YX.
+        n_rot : int
+            Number of 90 degree rotations.
+        axes : Tuple[int, int]
+            Axes along which to rotate the patch.
 
         Returns
         -------
-        Tuple[str]
-            Transform arguments.
+        np.ndarray
+            Transformed patch.
         """
-        return ("p", "is_3D")
+        return np.ascontiguousarray(np.rot90(patch, k=n_rot, axes=axes))

careamics/utils/__init__.py

@@ -1,6 +1,5 @@
 """Utils module."""
 
-
 __all__ = [
     "cwd",
     "get_ram_size",

careamics/utils/base_enum.py

@@ -1,9 +1,25 @@
+"""A base class for Enum that allows checking if a value is in the Enum."""
+
 from enum import Enum, EnumMeta
 from typing import Any
 
 
 class _ContainerEnum(EnumMeta):
+    """Metaclass for Enum with __contains__ method."""
+
     def __contains__(cls, item: Any) -> bool:
+        """Check if an item is in the Enum.
+
+        Parameters
+        ----------
+        item : Any
+            Item to check.
+
+        Returns
+        -------
+        bool
+            True if the item is in the Enum, False otherwise.
+        """
         try:
             cls(item)
         except ValueError:
@@ -12,6 +28,18 @@ class _ContainerEnum(EnumMeta):
 
     @classmethod
     def has_value(cls, value: Any) -> bool:
+        """Check if a value is in the Enum.
+
+        Parameters
+        ----------
+        value : Any
+            Value to check.
+
+        Returns
+        -------
+        bool
+            True if the value is in the Enum, False otherwise.
+        """
         return value in cls._value2member_map_
 
 

careamics/utils/context.py

@@ -3,6 +3,7 @@ Context submodule.
 
 A convenience function to change the working directory in order to save data.
 """
+
 import os
 from contextlib import contextmanager
 from pathlib import Path

careamics/utils/logging.py

@@ -3,6 +3,7 @@ Logging submodule.
 
 The methods are responsible for the in-console logger.
 """
+
 import logging
 import sys
 import time

careamics/utils/metrics.py

@@ -3,6 +3,7 @@ Metrics submodule.
 
 This module contains various metrics and a metrics tracking class.
 """
+
 from typing import Union
 
 import numpy as np

careamics/utils/path_utils.py

@@ -1,3 +1,5 @@
+"""Utility functions for paths."""
+
 from pathlib import Path
 from typing import Union
 

careamics/utils/ram.py

@@ -1,3 +1,5 @@
+"""Utility function to get RAM size."""
+
 import psutil
 
 

careamics/utils/receptive_field.py

@@ -1,102 +1,108 @@
 """Receptive field calculation for computing the tile overlap."""
 
+# TODO better docstring and function names
 # Adapted from: https://github.com/frgfm/torch-scan
 
-import math
-import warnings
-from typing import Tuple, Union
+# import math
+# import warnings
+# from typing import Tuple, Union
 
-from torch import Tensor, nn
-from torch.nn import Module
-from torch.nn.modules.batchnorm import _BatchNorm
-from torch.nn.modules.conv import _ConvNd, _ConvTransposeNd
-from torch.nn.modules.pooling import (
-    _AdaptiveAvgPoolNd,
-    _AdaptiveMaxPoolNd,
-    _AvgPoolNd,
-    _MaxPoolNd,
-)
+# from torch import Tensor, nn
+# from torch.nn import Module
+# from torch.nn.modules.batchnorm import _BatchNorm
+# from torch.nn.modules.conv import _ConvNd, _ConvTransposeNd
+# from torch.nn.modules.pooling import (
+#     _AdaptiveAvgPoolNd,
+#     _AdaptiveMaxPoolNd,
+#     _AvgPoolNd,
+#     _MaxPoolNd,
+# )
 
 
-def module_rf(module: Module, inp: Tensor, out: Tensor) -> Tuple[float, float, float]:
-    """Estimate the spatial receptive field of the module.
+# def module_rf(module: Module, inp: Tensor, out: Tensor) -> Tuple[float, float, float]:
+#     """Estimate the spatial receptive field of the module.
 
-    Args:
-        module (torch.nn.Module): PyTorch module
-        inp (torch.Tensor): input to the module
-        out (torch.Tensor): output of the module
-    Returns:
-        receptive field
-        effective stride
-        effective padding
-    """
-    if isinstance(
-        module,
-        (
-            nn.Identity,
-            nn.Flatten,
-            nn.ReLU,
-            nn.ELU,
-            nn.LeakyReLU,
-            nn.ReLU6,
-            nn.Tanh,
-            nn.Sigmoid,
-            _BatchNorm,
-            nn.Dropout,
-            nn.Linear,
-        ),
-    ):
-        return 1.0, 1.0, 0.0
-    elif isinstance(module, _ConvTransposeNd):
-        return rf_convtransposend(module, inp, out)
-    elif isinstance(module, (_ConvNd, _MaxPoolNd, _AvgPoolNd)):
-        return rf_aggregnd(module, inp, out)
-    elif isinstance(module, (_AdaptiveMaxPoolNd, _AdaptiveAvgPoolNd)):
-        return rf_adaptive_poolnd(module, inp, out)
-    else:
-        warnings.warn(
-            f"Module type not supported: {module.__class__.__name__}", stacklevel=1
-        )
-        return 1.0, 1.0, 0.0
+#     Parameters
+#     ----------
+#     module : Module
+#         Module to estimate the receptive field.
+#     inp : Tensor
+#         Input tensor.
+#     out : Tensor
+#         Output tensor.
 
+#     Returns
+#     -------
+#     Tuple[float, float, float]
+#         Receptive field, effective stride and padding.
+#     """
+#     if isinstance(
+#         module,
+#         (
+#             nn.Identity,
+#             nn.Flatten,
+#             nn.ReLU,
+#             nn.ELU,
+#             nn.LeakyReLU,
+#             nn.ReLU6,
+#             nn.Tanh,
+#             nn.Sigmoid,
+#             _BatchNorm,
+#             nn.Dropout,
+#             nn.Linear,
+#         ),
+#     ):
+#         return 1.0, 1.0, 0.0
+#     elif isinstance(module, _ConvTransposeNd):
+#         return rf_convtransposend(module, inp, out)
+#     elif isinstance(module, (_ConvNd, _MaxPoolNd, _AvgPoolNd)):
+#         return rf_aggregnd(module, inp, out)
+#     elif isinstance(module, (_AdaptiveMaxPoolNd, _AdaptiveAvgPoolNd)):
+#         return rf_adaptive_poolnd(module, inp, out)
+#     else:
+#         warnings.warn(
+#             f"Module type not supported: {module.__class__.__name__}", stacklevel=1
+#         )
+#         return 1.0, 1.0, 0.0
 
-def rf_convtransposend(
-    module: _ConvTransposeNd, _: Tensor, __: Tensor
-) -> Tuple[float, float, float]:
-    k = (
-        module.kernel_size[0]
-        if isinstance(module.kernel_size, tuple)
-        else module.kernel_size
-    )
-    s = module.stride[0] if isinstance(module.stride, tuple) else module.stride
-    return -k, 1.0 / s, 0.0
 
+# def rf_convtransposend(
+#     module: _ConvTransposeNd, _: Tensor, __: Tensor
+# ) -> Tuple[float, float, float]:
+#     k = (
+#         module.kernel_size[0]
+#         if isinstance(module.kernel_size, tuple)
+#         else module.kernel_size
+#     )
+#     s = module.stride[0] if isinstance(module.stride, tuple) else module.stride
+#     return -k, 1.0 / s, 0.0
 
-def rf_aggregnd(
-    module: Union[_ConvNd, _MaxPoolNd, _AvgPoolNd], _: Tensor, __: Tensor
-) -> Tuple[float, float, float]:
-    k = (
-        module.kernel_size[0]
-        if isinstance(module.kernel_size, tuple)
-        else module.kernel_size
-    )
-    if hasattr(module, "dilation"):
-        d = (
-            module.dilation[0]
-            if isinstance(module.dilation, tuple)
-            else module.dilation
-        )
-        k = d * (k - 1) + 1
-    s = module.stride[0] if isinstance(module.stride, tuple) else module.stride
-    p = module.padding[0] if isinstance(module.padding, tuple) else module.padding
-    return k, s, p  # type: ignore[return-value]
 
+# def rf_aggregnd(
+#     module: Union[_ConvNd, _MaxPoolNd, _AvgPoolNd], _: Tensor, __: Tensor
+# ) -> Tuple[float, float, float]:
+#     k = (
+#         module.kernel_size[0]
+#         if isinstance(module.kernel_size, tuple)
+#         else module.kernel_size
+#     )
+#     if hasattr(module, "dilation"):
+#         d = (
+#             module.dilation[0]
+#             if isinstance(module.dilation, tuple)
+#             else module.dilation
+#         )
+#         k = d * (k - 1) + 1
+#     s = module.stride[0] if isinstance(module.stride, tuple) else module.stride
+#     p = module.padding[0] if isinstance(module.padding, tuple) else module.padding
+#     return k, s, p  # type: ignore[return-value]
 
-def rf_adaptive_poolnd(
-    _: Union[_AdaptiveMaxPoolNd, _AdaptiveAvgPoolNd], inp: Tensor, out: Tensor
-) -> Tuple[int, int, float]:
-    stride = math.ceil(inp.shape[-1] / out.shape[-1])
-    kernel_size = stride
-    padding = (inp.shape[-1] - kernel_size * stride) / 2
 
-    return kernel_size, stride, padding
+# def rf_adaptive_poolnd(
+#     _: Union[_AdaptiveMaxPoolNd, _AdaptiveAvgPoolNd], inp: Tensor, out: Tensor
+# ) -> Tuple[int, int, float]:
+#     stride = math.ceil(inp.shape[-1] / out.shape[-1])
+#     kernel_size = stride
+#     padding = (inp.shape[-1] - kernel_size * stride) / 2
+
+#     return kernel_size, stride, padding

careamics/utils/torch_utils.py

@@ -3,6 +3,7 @@ Convenience functions using torch.
 
 These functions are used to control certain aspects and behaviours of PyTorch.
 """
+
 import inspect
 from typing import Dict, Union