pixel_manipulation

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

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