careamics 0.0.1__py3-none-any.whl → 0.0.3__py3-none-any.whl

careamics/transforms/pixel_manipulation.py

@@ -0,0 +1,407 @@
+"""
+Pixel manipulation methods.
+
+Pixel manipulation is used in N2V and similar algorithm to replace the value of
+masked pixels.
+"""
+
+from typing import Optional, Tuple
+
+import numpy as np
+
+from .struct_mask_parameters import StructMaskParameters
+
+
+def _apply_struct_mask(
+    patch: np.ndarray,
+    coords: np.ndarray,
+    struct_params: StructMaskParameters,
+    rng: Optional[np.random.Generator] = None,
+) -> np.ndarray:
+    """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
+    a random value.
+
+    Note that the structN2V mask is applied in 2D at the coordinates given by `coords`.
+
+    Parameters
+    ----------
+    patch : np.ndarray
+        Patch to be manipulated, 2D or 3D.
+    coords : np.ndarray
+        Coordinates of the ROI(subpatch) centers.
+    struct_params : StructMaskParameters
+        Parameters for the structN2V mask (axis and span).
+    rng : np.random.Generator or None
+        Random number generator.
+
+    Returns
+    -------
+    np.ndarray
+        Patch with the structN2V mask applied.
+    """
+    if rng is None:
+        rng = np.random.default_rng()
+
+    # relative axis
+    moving_axis = -1 - struct_params.axis
+
+    # Create a mask array
+    mask = np.expand_dims(
+        np.ones(struct_params.span), axis=list(range(len(patch.shape) - 1))
+    )  # (1, 1, span) or (1, span)
+
+    # Move the moving axis to the correct position
+    # i.e. the axis along which the coordinates should change
+    mask = np.moveaxis(mask, -1, moving_axis)
+    center = np.array(mask.shape) // 2
+
+    # Mark the center
+    mask[tuple(center.T)] = 0
+
+    # displacements from center
+    dx = np.indices(mask.shape)[:, mask == 1] - center[:, None]
+
+    # combine all coords (ndim, npts,) with all displacements (ncoords,ndim,)
+    mix = dx.T[..., None] + coords.T[None]
+    mix = mix.transpose([1, 0, 2]).reshape([mask.ndim, -1]).T
+
+    # delete entries that are out of bounds
+    mix = np.delete(mix, mix[:, moving_axis] < 0, axis=0)
+
+    max_bound = patch.shape[moving_axis] - 1
+    mix = np.delete(mix, mix[:, moving_axis] > max_bound, axis=0)
+
+    # replace neighbouring pixels with random values from flat dist
+    patch[tuple(mix.T)] = rng.uniform(patch.min(), patch.max(), size=mix.shape[0])
+
+    return patch
+
+
+def _odd_jitter_func(step: float, rng: np.random.Generator) -> np.ndarray:
+    """
+    Randomly sample a jitter to be applied to the masking grid.
+
+    This is done to account for cases where the step size is not an integer.
+
+    Parameters
+    ----------
+    step : float
+        Step size of the grid, output of np.linspace.
+    rng : np.random.Generator
+        Random number generator.
+
+    Returns
+    -------
+    np.ndarray
+        Array of random jitter to be added to the grid.
+    """
+    # Define the random jitter to be added to the grid
+    odd_jitter = np.where(np.floor(step) == step, 0, rng.integers(0, 2))
+
+    # Round the step size to the nearest integer depending on the jitter
+    return np.floor(step) if odd_jitter == 0 else np.ceil(step)
+
+
+def _get_stratified_coords(
+    mask_pixel_perc: float,
+    shape: Tuple[int, ...],
+    rng: Optional[np.random.Generator] = None,
+) -> np.ndarray:
+    """
+    Generate coordinates of the pixels to mask.
+
+    Randomly selects the coordinates of the pixels to mask in a stratified way, i.e.
+    the distance between masked pixels is approximately the same.
+
+    Parameters
+    ----------
+    mask_pixel_perc : float
+        Actual (quasi) percentage of masked pixels across the whole image. Used in
+        calculating the distance between masked pixels across each axis.
+    shape : Tuple[int, ...]
+        Shape of the input patch.
+    rng : np.random.Generator or None
+        Random number generator.
+
+    Returns
+    -------
+    np.ndarray
+        Array of coordinates of the masked pixels.
+    """
+    if len(shape) < 2 or len(shape) > 3:
+        raise ValueError(
+            "Calculating coordinates is only possible for 2D and 3D patches"
+        )
+
+    if rng is None:
+        rng = np.random.default_rng()
+
+    mask_pixel_distance = np.round((100 / mask_pixel_perc) ** (1 / len(shape))).astype(
+        np.int32
+    )
+
+    # Define a grid of coordinates for each axis in the input patch and the step size
+    pixel_coords = []
+    steps = []
+    for axis_size in shape:
+        # make sure axis size is evenly divisible by box size
+        num_pixels = int(np.ceil(axis_size / mask_pixel_distance))
+        axis_pixel_coords, step = np.linspace(
+            0, axis_size, num_pixels, dtype=np.int32, endpoint=False, retstep=True
+        )
+        # explain
+        pixel_coords.append(axis_pixel_coords.T)
+        steps.append(step)
+
+    # Create a meshgrid of coordinates for each axis in the input patch
+    coordinate_grid_list = np.meshgrid(*pixel_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)  # type: ignore
+        * np.ones_like(coordinate_grid).astype(np.int32)
+        - 1,
+        size=coordinate_grid.shape,
+        endpoint=True,
+    )
+    coordinate_grid += grid_random_increment
+    coordinate_grid = np.clip(coordinate_grid, 0, np.array(shape) - 1)
+    return coordinate_grid
+
+
+def _create_subpatch_center_mask(
+    subpatch: np.ndarray, center_coords: np.ndarray
+) -> np.ndarray:
+    """Create a mask with the center of the subpatch masked.
+
+    Parameters
+    ----------
+    subpatch : np.ndarray
+        Subpatch to be manipulated.
+    center_coords : np.ndarray
+        Coordinates of the original center before possible crop.
+
+    Returns
+    -------
+    np.ndarray
+        Mask with the center of the subpatch masked.
+    """
+    mask = np.ones(subpatch.shape)
+    mask[tuple(center_coords)] = 0
+    return np.ma.make_mask(mask)  # type: ignore
+
+
+def _create_subpatch_struct_mask(
+    subpatch: np.ndarray, center_coords: np.ndarray, struct_params: StructMaskParameters
+) -> np.ndarray:
+    """Create a structN2V mask for the subpatch.
+
+    Parameters
+    ----------
+    subpatch : np.ndarray
+        Subpatch to be manipulated.
+    center_coords : np.ndarray
+        Coordinates of the original center before possible crop.
+    struct_params : StructMaskParameters
+        Parameters for the structN2V mask (axis and span).
+
+    Returns
+    -------
+    np.ndarray
+        StructN2V mask for the subpatch.
+    """
+    # Create a mask with the center of the subpatch masked
+    mask_placeholder = np.ones(subpatch.shape)
+
+    # reshape to move the struct axis to the first position
+    mask_reshaped = np.moveaxis(mask_placeholder, struct_params.axis, 0)
+
+    # create the mask index for the struct axis
+    mask_index = slice(
+        max(0, center_coords.take(struct_params.axis) - (struct_params.span - 1) // 2),
+        min(
+            1 + center_coords.take(struct_params.axis) + (struct_params.span - 1) // 2,
+            subpatch.shape[struct_params.axis],
+        ),
+    )
+    mask_reshaped[struct_params.axis][mask_index] = 0
+
+    # reshape back to the original shape
+    mask = np.moveaxis(mask_reshaped, 0, struct_params.axis)
+
+    return np.ma.make_mask(mask)  # type: ignore
+
+
+def uniform_manipulate(
+    patch: np.ndarray,
+    mask_pixel_percentage: float,
+    subpatch_size: int = 11,
+    remove_center: bool = True,
+    struct_params: Optional[StructMaskParameters] = None,
+    rng: Optional[np.random.Generator] = None,
+) -> Tuple[np.ndarray, np.ndarray]:
+    """
+    Manipulate pixels by replacing them with a neighbor values.
+
+    Manipulated pixels are selected unformly selected in a subpatch, away from a grid
+    with an approximate uniform probability to be selected across the whole patch.
+    If `struct_params` is not None, an additional structN2V mask is applied to the
+    data, replacing the pixels in the mask with random values (excluding the pixel
+    already manipulated).
+
+    Parameters
+    ----------
+    patch : np.ndarray
+        Image patch, 2D or 3D, shape (y, x) or (z, y, x).
+    mask_pixel_percentage : float
+        Approximate percentage of pixels to be masked.
+    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.
+    struct_params : StructMaskParameters or None
+        Parameters for the structN2V mask (axis and span).
+    rng : np.random.Generator or None
+        Random number generator.
+
+    Returns
+    -------
+    Tuple[np.ndarray]
+        Tuple containing the manipulated patch and the corresponding mask.
+    """
+    if rng is None:
+        rng = np.random.default_rng()
+
+    # Get the coordinates of the pixels to be replaced
+    transformed_patch = patch.copy()
+
+    subpatch_centers = _get_stratified_coords(mask_pixel_percentage, patch.shape, rng)
+
+    # Generate coordinate grid for subpatch
+    roi_span_full = np.arange(
+        -np.floor(subpatch_size / 2), np.ceil(subpatch_size / 2)
+    ).astype(np.int32)
+
+    # Remove the center pixel from the grid if needed
+    roi_span = roi_span_full[roi_span_full != 0] if remove_center else roi_span_full
+
+    # Randomly select coordinates from the grid
+    random_increment = rng.choice(roi_span, size=subpatch_centers.shape)
+
+    # Clip the coordinates to the patch size
+    replacement_coords = np.clip(
+        subpatch_centers + random_increment,
+        0,
+        [patch.shape[i] - 1 for i in range(len(patch.shape))],
+    )
+
+    # Get the replacement pixels from all subpatchs
+    replacement_pixels = patch[tuple(replacement_coords.T.tolist())]
+
+    # Replace the original pixels with the replacement pixels
+    transformed_patch[tuple(subpatch_centers.T.tolist())] = replacement_pixels
+    mask = np.where(transformed_patch != patch, 1, 0).astype(np.uint8)
+
+    if struct_params is not None:
+        transformed_patch = _apply_struct_mask(
+            transformed_patch, subpatch_centers, struct_params
+        )
+
+    return (
+        transformed_patch,
+        mask,
+    )
+
+
+def median_manipulate(
+    patch: np.ndarray,
+    mask_pixel_percentage: float,
+    subpatch_size: int = 11,
+    struct_params: Optional[StructMaskParameters] = None,
+    rng: Optional[np.random.Generator] = None,
+) -> Tuple[np.ndarray, np.ndarray]:
+    """
+    Manipulate pixels by replacing them with the median of their surrounding subpatch.
+
+    N2V2 version, manipulated pixels are selected randomly away from a grid with an
+    approximate uniform probability to be selected across the whole patch.
+
+    If `struct_params` is not None, an additional structN2V mask is applied to the data,
+    replacing the pixels in the mask with random values (excluding the pixel already
+    manipulated).
+
+    Parameters
+    ----------
+    patch : np.ndarray
+        Image patch, 2D or 3D, shape (y, x) or (z, y, x).
+    mask_pixel_percentage : floar
+        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 : StructMaskParameters or None, optional
+        Parameters for the structN2V mask (axis and span).
+    rng : np.random.Generator or None, optional
+        Random number generato, by default None.
+
+    Returns
+    -------
+    Tuple[np.ndarray]
+           Tuple containing the manipulated patch, the original patch and the mask.
+    """
+    if rng is None:
+        rng = np.random.default_rng()
+
+    transformed_patch = patch.copy()
+
+    # Get the coordinates of the pixels to be replaced
+    subpatch_centers = _get_stratified_coords(mask_pixel_percentage, patch.shape, rng)
+
+    # Generate coordinate grid for subpatch
+    roi_span = np.array(
+        [-np.floor(subpatch_size / 2), np.ceil(subpatch_size / 2)]
+    ).astype(np.int32)
+
+    subpatch_crops_span_full = subpatch_centers[np.newaxis, ...].T + roi_span
+
+    # Dimensions n dims, n centers, (min, max)
+    subpatch_crops_span_clipped = np.clip(
+        subpatch_crops_span_full,
+        a_min=np.zeros_like(patch.shape)[:, np.newaxis, np.newaxis],
+        a_max=np.array(patch.shape)[:, np.newaxis, np.newaxis],
+    )
+
+    for idx in range(subpatch_crops_span_clipped.shape[1]):
+        subpatch_coords = subpatch_crops_span_clipped[:, idx, ...]
+        idxs = [
+            slice(x[0], x[1]) if x[1] - x[0] > 0 else slice(0, 1)
+            for x in subpatch_coords
+        ]
+        subpatch = patch[tuple(idxs)]
+        subpatch_center_adjusted = subpatch_centers[idx] - subpatch_coords[:, 0]
+
+        if struct_params is None:
+            subpatch_mask = _create_subpatch_center_mask(
+                subpatch, subpatch_center_adjusted
+            )
+        else:
+            subpatch_mask = _create_subpatch_struct_mask(
+                subpatch, subpatch_center_adjusted, struct_params
+            )
+        transformed_patch[tuple(subpatch_centers[idx])] = np.median(
+            subpatch[subpatch_mask]
+        )
+
+    mask = np.where(transformed_patch != patch, 1, 0).astype(np.uint8)
+
+    if struct_params is not None:
+        transformed_patch = _apply_struct_mask(
+            transformed_patch, subpatch_centers, struct_params
+        )
+
+    return (
+        transformed_patch,
+        mask,
+    )

careamics/transforms/struct_mask_parameters.py

@@ -0,0 +1,20 @@
+"""Class representing the parameters of structN2V masks."""
+
+from dataclasses import dataclass
+from typing import Literal
+
+
+@dataclass
+class StructMaskParameters:
+    """Parameters of structN2V masks.
+
+    Attributes
+    ----------
+    axis : Literal[0, 1]
+        Axis along which to apply the mask, horizontal (0) or vertical (1).
+    span : int
+        Span of the mask.
+    """
+
+    axis: Literal[0, 1]
+    span: int

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

@@ -0,0 +1,88 @@
+"""Test-time augmentations."""
+
+from torch import Tensor, flip, mean, rot90, stack
+
+
+class ImageRestorationTTA:
+    """
+    Test-time augmentation for image restoration tasks.
+
+    The augmentation is performed using all 90 deg rotations and their flipped version,
+    as well as the original image flipped.
+
+    Tensors should be of shape SC(Z)YX.
+
+    This transformation is used in the LightningModule in order to perform test-time
+    augmentation.
+    """
+
+    def forward(self, input_tensor: Tensor) -> list[Tensor]:
+        """
+        Apply test-time augmentation to the input tensor.
+
+        Parameters
+        ----------
+        input_tensor : Tensor
+            Input tensor, shape SC(Z)YX.
+
+        Returns
+        -------
+        list of torch.Tensor
+            List of augmented tensors.
+        """
+        # axes: only applies to YX axes
+        axes = (-2, -1)
+
+        augmented = [
+            # original
+            input_tensor,
+            # rotations
+            rot90(input_tensor, 1, dims=axes),
+            rot90(input_tensor, 2, dims=axes),
+            rot90(input_tensor, 3, dims=axes),
+            # original flipped
+            flip(input_tensor, dims=(axes[0],)),
+            flip(input_tensor, dims=(axes[1],)),
+        ]
+
+        # rotated once, flipped
+        augmented.extend(
+            [
+                flip(augmented[1], dims=(axes[0],)),
+                flip(augmented[1], dims=(axes[1],)),
+            ]
+        )
+
+        return augmented
+
+    def backward(self, x: list[Tensor]) -> Tensor:
+        """Undo the test-time augmentation.
+
+        Parameters
+        ----------
+        x : Any
+            List of augmented tensors of shape SC(Z)YX.
+
+        Returns
+        -------
+        Any
+            Original tensor.
+        """
+        axes = (-2, -1)
+
+        reverse = [
+            # original
+            x[0],
+            # rotated
+            rot90(x[1], -1, dims=axes),
+            rot90(x[2], -2, dims=axes),
+            rot90(x[3], -3, dims=axes),
+            # original flipped
+            flip(x[4], dims=(axes[0],)),
+            flip(x[5], dims=(axes[1],)),
+            # rotated once, flipped
+            rot90(flip(x[6], dims=(axes[0],)), -1, dims=axes),
+            rot90(flip(x[7], dims=(axes[1],)), -1, dims=axes),
+        ]
+
+        return mean(stack(reverse), dim=0)

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