careamics 0.0.19__py3-none-any.whl

careamics/transforms/pixel_manipulation_torch.py

@@ -0,0 +1,388 @@
+"""N2V manipulation functions for PyTorch."""
+
+import torch
+
+from .struct_mask_parameters import StructMaskParameters
+
+
+def _apply_struct_mask_torch(
+    patch: torch.Tensor,
+    coords: torch.Tensor,
+    struct_params: StructMaskParameters,
+    rng: torch.Generator | None = None,
+) -> torch.Tensor:
+    """Apply structN2V masks to patch.
+
+    Each point in `coords` corresponds to the center of a mask. Masks are parameterized
+    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 : torch.Tensor
+        Patch to be manipulated, (batch, y, x) or (batch, z, y, x).
+    coords : torch.Tensor
+        Coordinates of the ROI (subpatch) centers.
+    struct_params : StructMaskParameters
+        Parameters for the structN2V mask (axis and span).
+    rng : torch.Generator, optional
+        Random number generator.
+
+    Returns
+    -------
+    torch.Tensor
+        Patch with the structN2V mask applied.
+    """
+    if rng is None:
+        rng = torch.Generator(device=patch.device)
+
+    # Relative axis
+    moving_axis = -1 - struct_params.axis
+
+    # Create a mask array
+    mask_shape = [1] * len(patch.shape)
+    mask_shape[moving_axis] = struct_params.span
+    mask = torch.ones(mask_shape, device=patch.device)
+
+    center = torch.tensor(mask.shape, device=patch.device) // 2
+
+    # Mark the center
+    mask[tuple(center)] = 0
+
+    # Displacements from center
+    displacements = torch.stack(torch.where(mask == 1)) - center.unsqueeze(1)
+
+    # Combine all coords (ndim, npts) with all displacements (ncoords, ndim)
+    mix = displacements.T.unsqueeze(-1) + coords.T.unsqueeze(0)
+    mix = mix.permute([1, 0, 2]).reshape([mask.ndim, -1]).T
+
+    # Filter out invalid indices
+    valid_indices = (mix[:, moving_axis] >= 0) & (
+        mix[:, moving_axis] < patch.shape[moving_axis]
+    )
+    mix = mix[valid_indices]
+
+    mins = patch.min(-1)[0].min(-1)[0]
+    maxs = patch.max(-1)[0].max(-1)[0]
+    for i in range(patch.shape[0]):
+        batch_coords = mix[mix[:, 0] == i]
+        min_ = mins[i].item()
+        max_ = maxs[i].item()
+        random_values = torch.empty(len(batch_coords), device=patch.device).uniform_(
+            min_, max_, generator=rng
+        )
+        patch[tuple(batch_coords[:, i] for i in range(patch.ndim))] = random_values
+
+    return patch
+
+
+def _get_stratified_coords_torch(
+    mask_pixel_perc: float,
+    shape: tuple[int, ...],
+    rng: torch.Generator,
+) -> torch.Tensor:
+    """
+    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. This is achieved by
+    defining a grid and sampling a pixel in each grid square. The grid is defined such
+    that the resulting density of masked pixels is the desired masked pixel percentage.
+
+    Parameters
+    ----------
+    mask_pixel_perc : float
+        Expected value for percentage of masked pixels across the whole image.
+    shape : tuple[int, ...]
+        Shape of the input patch.
+    rng : torch.Generator or None
+        Random number generator.
+
+    Returns
+    -------
+    np.ndarray
+        Array of coordinates of the masked pixels.
+    """
+    # Implementation logic:
+    #    find a box size s.t sampling 1 pixel within the box will result in the desired
+    # pixel percentage. Make a grid of these boxes that cover the patch (the area of
+    # the grid will be greater than or equal to the area of the patch) and sample 1
+    # pixel in each box. The density of masked pixels is an intensive property therefore
+    # any subset of this area will have the desired expected masked pixel percentage.
+    # We can get our desired patch with our desired expected masked pixel percentage by
+    # simply filtering out masked pixels that lie outside of our patch bounds.
+
+    batch_size = shape[0]
+    spatial_shape = shape[1:]
+
+    n_dims = len(spatial_shape)
+    expected_area_per_pixel = 1 / (mask_pixel_perc / 100)
+
+    # keep the grid size in floats for a more accurate expected masked pixel percentage
+    grid_size = expected_area_per_pixel ** (1 / n_dims)
+    grid_dims = torch.ceil(torch.tensor(spatial_shape) / grid_size).int()
+
+    # coords on a fixed grid (top left corner)
+    coords = torch.stack(
+        torch.meshgrid(
+            torch.arange(batch_size, dtype=torch.float),
+            *[torch.arange(0, grid_dims[i].item()) * grid_size for i in range(n_dims)],
+            indexing="ij",
+        ),
+        -1,
+    ).reshape(-1, n_dims + 1)
+
+    # add random offset to get a random coord in each grid box
+    # also keep the offset in floats
+    offset = (
+        torch.rand((len(coords), n_dims), device=rng.device, generator=rng) * grid_size
+    )
+    coords = coords.to(rng.device)
+    coords[:, 1:] += offset
+    coords = torch.floor(coords).int()
+
+    # filter pixels out of bounds
+    out_of_bounds = (
+        coords[:, 1:]
+        >= torch.tensor(spatial_shape, device=rng.device).reshape(1, n_dims)
+    ).any(1)
+    coords = coords[~out_of_bounds]
+    return coords
+
+
+def uniform_manipulate_torch(
+    patch: torch.Tensor,
+    mask_pixel_percentage: float,
+    subpatch_size: int = 11,
+    remove_center: bool = True,
+    struct_params: StructMaskParameters | None = None,
+    rng: torch.Generator | None = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """
+    Manipulate pixels by replacing them with a neighbor values.
+
+    # TODO add more details, especially about batch
+
+    Manipulated pixels are selected uniformly 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 : torch.Tensor
+        Image patch, 2D or 3D, shape (y, x) or (z, y, x). # TODO batch and channel.
+    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 : torch.default_generator or None
+        Random number generator.
+
+    Returns
+    -------
+    tuple[torch.Tensor, torch.Tensor]
+        tuple containing the manipulated patch and the corresponding mask.
+    """
+    if rng is None:
+        rng = torch.Generator(device=patch.device)
+        # TODO do we need seed ?
+
+    # create a copy of the patch
+    transformed_patch = patch.clone()
+
+    # get the coordinates of the pixels to be masked
+    subpatch_centers = _get_stratified_coords_torch(
+        mask_pixel_percentage, patch.shape, rng
+    )
+    subpatch_centers = subpatch_centers.to(device=patch.device)
+
+    # TODO refactor with non negative indices?
+    # arrange the list of indices to represent the ROI around the pixel to be masked
+    roi_span_full = torch.arange(
+        -(subpatch_size // 2),
+        subpatch_size // 2 + 1,
+        dtype=torch.int32,
+        device=patch.device,
+    )
+
+    # remove the center pixel from the ROI
+    roi_span = roi_span_full[roi_span_full != 0] if remove_center else roi_span_full
+
+    # create a random increment to select the replacement value
+    # this increment is added to the center coordinates
+    random_increment = roi_span[
+        torch.randint(
+            low=min(roi_span),
+            high=max(roi_span) + 1,
+            # one less coord dim: we shouldn't add a random increment to the batch coord
+            size=(subpatch_centers.shape[0], subpatch_centers.shape[1] - 1),
+            generator=rng,
+            device=patch.device,
+        )
+    ]
+
+    # compute the replacement pixel coordinates
+    replacement_coords = subpatch_centers.clone()
+    # only add random increment to the spatial dimensions, not the batch dimension
+    replacement_coords[:, 1:] = torch.clamp(
+        replacement_coords[:, 1:] + random_increment,
+        torch.zeros_like(torch.tensor(patch.shape[1:])).to(device=patch.device),
+        torch.tensor([v - 1 for v in patch.shape[1:]]).to(device=patch.device),
+    )
+
+    # replace the pixels in the patch
+    # tuples and transpose are needed for proper indexing
+    replacement_pixels = patch[tuple(replacement_coords.T)]
+    transformed_patch[tuple(subpatch_centers.T)] = replacement_pixels
+
+    # create a mask representing the masked pixels
+    mask = (transformed_patch != patch).to(dtype=torch.uint8)
+
+    # apply structN2V mask if needed
+    if struct_params is not None:
+        transformed_patch = _apply_struct_mask_torch(
+            transformed_patch, subpatch_centers, struct_params, rng
+        )
+
+    return transformed_patch, mask
+
+
+def median_manipulate_torch(
+    batch: torch.Tensor,
+    mask_pixel_percentage: float,
+    subpatch_size: int = 11,
+    struct_params: StructMaskParameters | None = None,
+    rng: torch.Generator | None = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """
+    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
+    ----------
+    batch : torch.Tensor
+        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.
+    struct_params : StructMaskParameters or None, optional
+        Parameters for the structN2V mask (axis and span).
+    rng : torch.default_generator or None, optional
+        Random number generator, by default None.
+
+    Returns
+    -------
+    tuple[torch.Tensor, torch.Tensor, torch.Tensor]
+           tuple containing the manipulated patch, the original patch and the mask.
+    """
+    # get the coordinates of the future ROI centers
+    subpatch_center_coordinates = _get_stratified_coords_torch(
+        mask_pixel_percentage, batch.shape, rng
+    ).to(
+        device=batch.device
+    )  # (num_coordinates, batch + num_spatial_dims)
+
+    # Calculate the padding value for the input tensor
+    pad_value = subpatch_size // 2
+
+    # Generate all offsets for the ROIs. Iteration starting from 1 to skip the batch
+    offsets = torch.meshgrid(
+        [
+            torch.arange(-pad_value, pad_value + 1, device=batch.device)
+            for _ in range(1, subpatch_center_coordinates.shape[1])
+        ],
+        indexing="ij",
+    )
+    offsets = torch.stack(
+        [axis_offset.flatten() for axis_offset in offsets], dim=1
+    )  # (subpatch_size**2, num_spatial_dims)
+
+    # Create the list to assemble coordinates of the ROIs centers for each axis
+    coords_axes = []
+    # Create the list to assemble the span of coordinates defining the ROIs for each
+    # axis
+    coords_expands = []
+    for d in range(subpatch_center_coordinates.shape[1]):
+        coords_axes.append(subpatch_center_coordinates[:, d])
+        if d == 0:
+            # For batch dimension coordinates are not expanded (no offsets)
+            coords_expands.append(
+                subpatch_center_coordinates[:, d]
+                .unsqueeze(1)
+                .expand(-1, subpatch_size ** offsets.shape[1])
+            )  # (num_coordinates, subpatch_size**num_spacial_dims)
+        else:
+            # For spatial dimensions, coordinates are expanded with offsets, creating
+            # spans
+            coords_expands.append(
+                (
+                    subpatch_center_coordinates[:, d].unsqueeze(1) + offsets[:, d - 1]
+                ).clamp(0, batch.shape[d] - 1)
+            )  # (num_coordinates, subpatch_size**num_spacial_dims)
+
+    # create array of rois by indexing the batch with gathered coordinates
+    rois = batch[
+        tuple(coords_expands)
+    ]  # (num_coordinates, subpatch_size**num_spacial_dims)
+
+    if struct_params is not None:
+        # Create the structN2V mask
+        h, w = torch.meshgrid(
+            torch.arange(subpatch_size), torch.arange(subpatch_size), indexing="ij"
+        )
+        center_idx = subpatch_size // 2
+        halfspan = (struct_params.span - 1) // 2
+
+        # Determine the axis along which to apply the mask
+        if struct_params.axis == 0:
+            center_axis = h
+            span_axis = w
+        else:
+            center_axis = w
+            span_axis = h
+
+        # Create the mask
+        struct_mask = (
+            ~(
+                (center_axis == center_idx)
+                & (span_axis >= center_idx - halfspan)
+                & (span_axis <= center_idx + halfspan)
+            )
+        ).flatten()
+        rois_filtered = rois[:, struct_mask]
+    else:
+        # Remove the center pixel value from the rois
+        center_idx = (subpatch_size ** offsets.shape[1]) // 2
+        rois_filtered = torch.cat(
+            [rois[:, :center_idx], rois[:, center_idx + 1 :]], dim=1
+        )
+
+    # compute the medians.
+    medians = rois_filtered.median(dim=1).values  # (num_coordinates,)
+
+    # Update the output tensor with medians
+    output_batch = batch.clone()
+    output_batch[tuple(coords_axes)] = medians
+    mask = torch.where(output_batch != batch, 1, 0).to(torch.uint8)
+
+    if struct_params is not None:
+        output_batch = _apply_struct_mask_torch(
+            output_batch, subpatch_center_coordinates, struct_params
+        )
+
+    return output_batch, 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,131 @@
+"""XY flip transform."""
+
+import numpy as np
+from numpy.typing import NDArray
+
+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: int | None = 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: 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 : np.ndarray
+            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
+        -------
+        Tuple[np.ndarray, Optional[np.ndarray]]
+            Transformed patch and target.
+        """
+        if self.rng.random() > self.p:
+            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, additional_transformed
+
+    def _apply(self, patch: NDArray, axis: int) -> 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))