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

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/pixel_manipulation.py

@@ -161,7 +161,7 @@ def _get_stratified_coords(
     coordinate_grid = np.array(coordinate_grid_list).reshape(len(shape), -1).T
 
     grid_random_increment = rng.integers(
-        _odd_jitter_func(float(max(steps)), rng)
+        _odd_jitter_func(float(max(steps)), rng)  # type: ignore
         * np.ones_like(coordinate_grid).astype(np.int32)
         - 1,
         size=coordinate_grid.shape,

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