Image Processing Module

@njit(parallel=True)
def downsample_axis(image: ArrayLike, level: int = 2, axis: int = 0) -> ArrayLike:
    """Numba accelerated downsampling for 3D images along a specified axis.

    Parameters
    ----------
    image: ArrayLike
        3D image to be downsampled.
    level: int
        Amount of downsampling.
    axis: int
        Axis along which to downsample (0, 1, or 2).

    Returns
    -------
    downsampled_image: ArrayLike
        3D downsampled image.

    """
    if axis == 0:
        new_length = image.shape[0] // level + (1 if image.shape[0] % level != 0 else 0)
        downsampled_image = np.zeros(
            (new_length, image.shape[1], image.shape[2]), dtype=image.dtype
        )

        for y in prange(image.shape[1]):
            for x in range(image.shape[2]):
                for z in range(new_length):
                    sum_value = 0.0
                    count = 0
                    for j in range(level):
                        original_index = z * level + j
                        if original_index < image.shape[0]:
                            sum_value += image[original_index, y, x]
                            count += 1
                    if count > 0:
                        downsampled_image[z, y, x] = sum_value / count

    elif axis == 1:
        new_length = image.shape[1] // level + (1 if image.shape[1] % level != 0 else 0)
        downsampled_image = np.zeros(
            (image.shape[0], new_length, image.shape[2]), dtype=image.dtype
        )

        for z in prange(image.shape[0]):
            for x in range(image.shape[2]):
                for y in range(new_length):
                    sum_value = 0.0
                    count = 0
                    for j in range(level):
                        original_index = y * level + j
                        if original_index < image.shape[1]:
                            sum_value += image[z, original_index, x]
                            count += 1
                    if count > 0:
                        downsampled_image[z, y, x] = sum_value / count

    elif axis == 2:
        new_length = image.shape[2] // level + (1 if image.shape[2] % level != 0 else 0)
        downsampled_image = np.zeros(
            (image.shape[0], image.shape[1], new_length), dtype=image.dtype
        )

        for z in prange(image.shape[0]):
            for y in range(image.shape[1]):
                for x in range(new_length):
                    sum_value = 0.0
                    count = 0
                    for j in range(level):
                        original_index = x * level + j
                        if original_index < image.shape[2]:
                            sum_value += image[z, y, original_index]
                            count += 1
                    if count > 0:
                        downsampled_image[z, y, x] = sum_value / count

    return downsampled_image

def downsample_image_anisotropic(
    image: ArrayLike, level: tuple[int, int, int] = (2, 6, 6)
) -> ArrayLike:
    """Numba accelerated anisotropic downsampling.

    Parameters
    ----------
    image: ArrayLike
        3D image to be downsampled
    level: tuple[int,int,int], default=(2,6,6)
        anisotropic downsampling level

    Returns
    -------
    downsampled_image: ArrayLike
        downsampled 3D image
    """
    downsampled_image = downsample_axis(
        downsample_axis(downsample_axis(image, level[0], 0), level[1], 1), level[2], 2
    )

    return downsampled_image

def estimate_shading(images: "Iterable[Future]") -> ArrayLike:
    """Estimate shading using stack of images and BaSiCPy.

    Parameters
    ----------
    images: Iterable[Future]
        Read futures yielding ZYX images. Pass a generator to load one stack at
        a time. Only YX maximum projections are retained for fitting.

    Returns
    -------
    shading_image: ArrayLike
        estimated shading image
    """
    # GPU

    import cupy as cp  # type: ignore
    from basicpy import BaSiC  # type: ignore

    maxz_images = []
    for image in images:
        maxz_images.append(cp.squeeze(cp.max(image.result(), axis=0)))
        del image

    maxz_images = cp.asnumpy(maxz_images).astype(np.uint16)
    gc.collect()
    cp.cuda.Stream.null.synchronize()
    cp.get_default_memory_pool().free_all_blocks()
    cp.get_default_pinned_memory_pool().free_all_blocks()

    working_size = [max(1, size // 2) for size in maxz_images.shape[-2:]]
    basic = BaSiC(
        get_darkfield=False,
        sort_intensity=True,
        working_size=working_size,
    )
    with redirect_stdout(io.StringIO()):
        basic.autotune(maxz_images[:])
        basic.fit(maxz_images[:])
    shading_correction = basic.flatfield.astype(np.float32) / np.max(
        basic.flatfield.astype(np.float32), axis=(0, 1)
    )

    del basic
    gc.collect()

    cp.cuda.Stream.null.synchronize()
    cp.get_default_memory_pool().free_all_blocks()
    cp.get_default_pinned_memory_pool().free_all_blocks()

    return shading_correction

def image_has_signal(
    image: ArrayLike,
    *,
    mask: ArrayLike | None = None,
    background_sigma: float = 5.0,
    minimum_pixels: int = 9,
) -> bool:
    """Screen YX or ZYX data for connected signal above a robust background.

    Median-filter the Z maximum projection over 3x3 pixels to suppress isolated
    hot pixels. Require a connected region of at least ``minimum_pixels`` above
    median + ``background_sigma`` times the MAD-based noise estimate (1.4826 MAD).
    An optional YX mask restricts both background estimation and accepted signal.
    This is an intensity-contrast screen, not a calibrated significance test.
    Empty images/masks return False; invalid dimensions or nonfinite data raise.
    """
    from scipy.ndimage import label, median_filter

    image = np.asarray(image)
    if image.ndim not in (2, 3):
        raise ValueError("Signal screening requires a YX or ZYX image.")
    if image.size == 0:
        return False
    projection = image.max(axis=0) if image.ndim == 3 else image
    if not np.all(np.isfinite(projection)):
        raise ValueError("Signal screening requires finite image values.")
    projection = median_filter(projection.astype(np.float32), size=3)
    pixels = projection if mask is None else projection[np.asarray(mask, dtype=bool)]
    if pixels.size == 0:
        return False
    background = np.median(pixels)
    noise = 1.4826 * np.median(np.abs(pixels - background))
    foreground = projection > background + background_sigma * noise
    if mask is not None:
        foreground &= np.asarray(mask, dtype=bool)
    labels, _ = label(foreground)
    return bool(np.any(np.bincount(labels.ravel())[1:] >= minimum_pixels))

def replace_hot_pixels(
    noise_map: ArrayLike, data: ArrayLike, threshold: float = 375.0
) -> ArrayLike:
    """Replace hot pixels with median values surrounding them.

    Parameters
    ----------
    noise_map: ArrayLike
        darkfield image collected at long exposure time to get hot pixels
    data: ArrayLike
        ND data [broadcast_dim,z,y,x]
    threshold: float, default = 375.0
        threshold to determine hot pixel

    Returns
    -------
    data: ArrayLike
        hotpixel corrected data
    """
    # GPU
    import cupy as cp  # type: ignore
    from cupyx.scipy import ndimage  # type: ignore

    data = cp.asarray(data, dtype=cp.float32)
    noise_map = cp.asarray(noise_map, dtype=cp.float32)

    # threshold darkfield_image to generate bad pixel matrix
    hot_pixels = cp.squeeze(cp.asarray(noise_map))
    hot_pixels[hot_pixels <= threshold] = 0
    hot_pixels[hot_pixels > threshold] = 1
    hot_pixels = hot_pixels.astype(cp.float32)
    inverted_hot_pixels = cp.ones_like(hot_pixels) - hot_pixels.copy()

    data = cp.asarray(data, dtype=cp.float32)
    for z_idx in range(data.shape[0]):
        median = ndimage.median_filter(data[z_idx, :, :], size=3)
        data[z_idx, :] = inverted_hot_pixels * data[z_idx, :] + hot_pixels * median

    data[data < 0] = 0

    data = cp.asnumpy(data).astype(np.uint16)
    gc.collect()
    cp.cuda.Stream.null.synchronize()
    cp.get_default_memory_pool().free_all_blocks()
    cp.get_default_pinned_memory_pool().free_all_blocks()

    return data