careamics 0.0.19__py3-none-any.whl

careamics/transforms/n2v_manipulate.py

@@ -0,0 +1,150 @@
+"""N2V manipulation transform."""
+
+from typing import Any, Literal
+
+import numpy as np
+from numpy.typing import NDArray
+
+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(Transform):
+    """
+    Default augmentation for the N2V model.
+
+    This transform expects C(Z)YX dimensions.
+
+    Parameters
+    ----------
+    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 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__(
+        self,
+        roi_size: int = 11,
+        masked_pixel_percentage: float = 0.2,
+        strategy: Literal[
+            "uniform", "median"
+        ] = SupportedPixelManipulation.UNIFORM.value,
+        remove_center: bool = True,
+        struct_mask_axis: Literal["horizontal", "vertical", "none"] = "none",
+        struct_mask_span: int = 5,
+        seed: int | None = None,
+    ):
+        """Constructor.
+
+        Parameters
+        ----------
+        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.
+        """
+        self.masked_pixel_percentage = masked_pixel_percentage
+        self.roi_size = roi_size
+        self.strategy = strategy
+        self.remove_center = remove_center  # TODO is this ever used?
+
+        if struct_mask_axis == SupportedStructAxis.NONE:
+            self.struct_mask: StructMaskParameters | None = None
+        else:
+            self.struct_mask = StructMaskParameters(
+                axis=0 if struct_mask_axis == SupportedStructAxis.HORIZONTAL else 1,
+                span=struct_mask_span,
+            )
+
+        # numpy random generator
+        self.rng = np.random.default_rng(seed=seed)
+
+    def __call__(
+        self, patch: NDArray, *args: Any, **kwargs: Any
+    ) -> tuple[NDArray, NDArray, NDArray]:
+        """Apply the transform to the image.
+
+        Parameters
+        ----------
+        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[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,
+                    struct_params=self.struct_mask,
+                    rng=self.rng,
+                )
+        elif self.strategy == SupportedPixelManipulation.MEDIAN:
+            # Iterate over the channels to apply manipulation separately
+            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,
+                    rng=self.rng,
+                )
+        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/n2v_manipulate_torch.py

@@ -0,0 +1,149 @@
+"""N2V manipulation transform for PyTorch."""
+
+import platform
+from typing import Any
+
+import torch
+
+from careamics.config.support import SupportedPixelManipulation, SupportedStructAxis
+from careamics.config.transformations import N2VManipulateConfig
+
+from .pixel_manipulation_torch import (
+    median_manipulate_torch,
+    uniform_manipulate_torch,
+)
+from .struct_mask_parameters import StructMaskParameters
+
+
+class N2VManipulateTorch:
+    """
+    Default augmentation for the N2V model.
+
+    This transform expects C(Z)YX dimensions.
+
+    Parameters
+    ----------
+    n2v_manipulate_config : N2VManipulateConfig
+        N2V manipulation configuration.
+    seed : Optional[int], optional
+        Random seed, by default None.
+    device : str
+        The device on which operations take place, e.g. "cuda", "cpu" or "mps".
+
+    Attributes
+    ----------
+    masked_pixel_percentage : float
+        Percentage of pixels to mask.
+    roi_size : int
+        Size of the replacement area.
+    strategy : Literal[ "uniform", "median" ]
+        Replacement 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__(
+        self,
+        n2v_manipulate_config: N2VManipulateConfig,
+        seed: int | None = None,
+        device: str | None = None,
+    ):
+        """Constructor.
+
+        Parameters
+        ----------
+        n2v_manipulate_config : N2VManipulateConfig
+            N2V manipulation configuration.
+        seed : Optional[int], optional
+            Random seed, by default None.
+        device : str
+            The device on which operations take place, e.g. "cuda", "cpu" or "mps".
+        """
+        self.masked_pixel_percentage = n2v_manipulate_config.masked_pixel_percentage
+        self.roi_size = n2v_manipulate_config.roi_size
+        self.strategy = n2v_manipulate_config.strategy
+        self.remove_center = n2v_manipulate_config.remove_center
+
+        if n2v_manipulate_config.struct_mask_axis == SupportedStructAxis.NONE:
+            self.struct_mask: StructMaskParameters | None = None
+        else:
+            self.struct_mask = StructMaskParameters(
+                axis=(
+                    0
+                    if n2v_manipulate_config.struct_mask_axis
+                    == SupportedStructAxis.HORIZONTAL
+                    else 1
+                ),
+                span=n2v_manipulate_config.struct_mask_span,
+            )
+
+        # PyTorch random generator
+        # TODO refactor into careamics.utils.torch_utils.get_device
+        if device is None:
+            if torch.cuda.is_available():
+                device = "cuda"
+            elif torch.backends.mps.is_available() and platform.processor() in (
+                "arm",
+                "arm64",
+            ):
+                device = "mps"
+            else:
+                device = "cpu"
+
+        self.rng = (
+            torch.Generator(device=device).manual_seed(seed)
+            if seed is not None
+            else torch.Generator(device=device)
+        )
+
+    def __call__(
+        self, batch: torch.Tensor, *args: Any, **kwargs: Any
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        """Apply the transform to the image.
+
+        Parameters
+        ----------
+        batch : torch.Tensor
+            Batch if image patches, 2D or 3D, shape BC(Z)YX.
+        *args : Any
+            Additional arguments, unused.
+        **kwargs : Any
+            Additional keyword arguments, unused.
+
+        Returns
+        -------
+        tuple[torch.Tensor, torch.Tensor, torch.Tensor]
+            Masked patch, original patch, and mask.
+        """
+        masked = torch.zeros_like(batch)
+        mask = torch.zeros_like(batch, dtype=torch.uint8)
+
+        if self.strategy == SupportedPixelManipulation.UNIFORM:
+            # Iterate over the channels to apply manipulation separately
+            for c in range(batch.shape[1]):
+                masked[:, c, ...], mask[:, c, ...] = uniform_manipulate_torch(
+                    patch=batch[:, c, ...],
+                    mask_pixel_percentage=self.masked_pixel_percentage,
+                    subpatch_size=self.roi_size,
+                    remove_center=self.remove_center,
+                    struct_params=self.struct_mask,
+                    rng=self.rng,
+                )
+        elif self.strategy == SupportedPixelManipulation.MEDIAN:
+            # Iterate over the channels to apply manipulation separately
+            for c in range(batch.shape[1]):
+                masked[:, c, ...], mask[:, c, ...] = median_manipulate_torch(
+                    batch=batch[:, c, ...],
+                    mask_pixel_percentage=self.masked_pixel_percentage,
+                    subpatch_size=self.roi_size,
+                    struct_params=self.struct_mask,
+                    rng=self.rng,
+                )
+        else:
+            raise ValueError(f"Unknown masking strategy ({self.strategy}).")
+
+        return masked, batch, mask

careamics/transforms/normalize.py

@@ -0,0 +1,374 @@
+"""Normalization and denormalization transforms for image patches."""
+
+import numpy as np
+import torch
+from numpy.typing import NDArray
+from torch import Tensor
+
+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))]
+
+
+def _reshape_stats_torch(stats: list[float], ndim: int) -> Tensor:
+    """Torch equivalent of `_reshape_stats` for broadcasting over image dims.
+
+    Parameters
+    ----------
+    stats : list of float
+        List of stats, mean or standard deviation.
+    ndim : int
+        Number of dimensions of the tensor, including the C channel.
+
+    Returns
+    -------
+    Tensor
+        Reshaped stats tensor.
+    """
+    t = torch.tensor(stats)
+    # Add singleton dimensions to match input tensor ndim for broadcasting
+    return t[(..., *[None] * (ndim - 1))]
+
+
+class Normalize(Transform):
+    """
+    Normalize an image or image patch.
+
+    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
+    ----------
+    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
+    ----------
+    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,
+        image_means: list[float],
+        image_stds: list[float],
+        target_means: list[float] | None = None,
+        target_stds: list[float] | None = None,
+    ):
+        """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: NDArray | None = None,
+        **additional_arrays: NDArray,
+    ) -> tuple[NDArray, NDArray | None, dict[str, NDArray]]:
+        """Apply the transform to the source patch and the target (optional).
+
+        Parameters
+        ----------
+        patch : NDArray
+            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
+        -------
+        tuple of NDArray
+            Transformed patch and target, the target can be returned as `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."
+            )
+        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)
+        stds = _reshape_stats(self.image_stds, patch.ndim)
+        norm_patch = self._apply(patch, means, stds)
+
+        # same for the target patch
+        if target is None:
+            norm_target = None
+        else:
+            if not self.target_means or not self.target_stds:
+                raise ValueError(
+                    "Target means and standard deviations must be provided "
+                    "if target is not None."
+                )
+            if len(self.target_means) == 0 and len(self.target_stds) == 0:
+                raise ValueError(
+                    "Target means and standard deviations must be provided "
+                    "if target is not None."
+                )
+            if len(self.target_means) != target.shape[0]:
+                raise ValueError(
+                    "Target means and standard deviations must have the same length "
+                    "as the target."
+                )
+            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)
+
+        return norm_patch, norm_target, additional_arrays
+
+    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.
+
+    Denormalization is performed expecting a zero mean and unit variance input. This
+    transform expects C(Z)YX dimensions.
+
+    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.
+
+    Parameters
+    ----------
+    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,
+        image_means: list[float],
+        image_stds: list[float],
+    ):
+        """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: NDArray) -> NDArray:
+        """Reverse the normalization operation for a batch of patches.
+
+        Parameters
+        ----------
+        patch : NDArray
+            Patch, 2D or 3D, shape BC(Z)YX.
+
+        Returns
+        -------
+        NDArray
+            Transformed array.
+        """
+        # 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."
+        #     )
+        # TODO for pn2v channel handling needs to be changed
+        means = _reshape_stats(self.image_means, patch.ndim)
+        stds = _reshape_stats(self.image_stds, patch.ndim)
+
+        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
+        ----------
+        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 array * (std + self.eps) + mean
+
+
+class TrainDenormalize:
+    """
+    Denormalize an image tensor for training-time tensors.
+
+    This class mirrors `Denormalize` but operates on torch tensors. It expects
+    the input tensor to have shape BC(Z)YX with the channel dimension at index 1.
+
+    Parameters
+    ----------
+    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,
+        image_means: list[float],
+        image_stds: list[float],
+    ) -> None:
+        """Initialize Denormalize transform.
+
+        Parameters
+        ----------
+        image_means : list of float
+            Mean values per channel.
+        image_stds : list of float
+            Standard deviation values per channel.
+        """
+        self.image_means = image_means
+        self.image_stds = image_stds
+        self.eps = 1e-6
+
+    def __call__(self, patch: Tensor) -> Tensor:
+        """Reverse the normalization operation for a batch of patches.
+
+        Parameters
+        ----------
+        patch : Tensor
+            Patch, 2D or 3D, shape BC(Z)YX.
+
+        Returns
+        -------
+        Tensor
+            Denormalized tensor with dtype float32.
+        """
+        # 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 {tuple(patch.shape)} for BC(Z)YX) "
+        #         f"don't match."
+        #     )
+        # TODO for pn2v channel handling needs to be changed
+
+        means = _reshape_stats_torch(self.image_means, patch.ndim).to(
+            device=patch.device, dtype=patch.dtype
+        )
+        stds = _reshape_stats_torch(self.image_stds, patch.ndim).to(
+            device=patch.device, dtype=patch.dtype
+        )
+
+        denorm_tensor = self._apply(
+            patch,
+            torch.swapaxes(means, 0, 1),  # swap axes as C channel is axis 1
+            torch.swapaxes(stds, 0, 1),
+        )
+
+        return denorm_tensor.float()
+
+    def _apply(self, array: Tensor, mean: Tensor, std: Tensor) -> Tensor:
+        """Apply the denormalization to the tensor.
+
+        Parameters
+        ----------
+        array : Tensor
+            Input tensor.
+        mean : Tensor
+            Mean values.
+        std : Tensor
+            Standard deviation values.
+
+        Returns
+        -------
+        Tensor
+            Denormalized tensor.
+        """
+        return array * (std + self.eps) + mean