deepivus/processing/preprocessing.py

"""Frame preprocessing utilities."""

import numpy as np


def ensure_uint8(images: np.ndarray) -> np.ndarray:
    """Convert any numeric image stack to uint8 with min-max scaling."""
    if images.dtype == np.uint8:
        return images
    images = images.astype(np.float32)
    min_val = float(images.min())
    max_val = float(images.max())
    if max_val <= min_val:
        return np.zeros(images.shape, dtype=np.uint8)
    scaled = (images - min_val) * (255.0 / (max_val - min_val))
    return scaled.astype(np.uint8)


def apply_center_black_circle(
    images: np.ndarray,
    diameter: int = 30,
    center_scale: float = 0.0,
    edge_scale: float = 0.5,
) -> np.ndarray:
    """Apply a soft radial attenuation mask centered in each frame.

    The attenuation is linear in radius from ``center_scale`` (at center)
    to ``edge_scale`` (at circle boundary), with no attenuation outside.
    """
    if images.ndim != 3:
        raise ValueError(f"Expected images shape (frames, rows, cols), got {images.shape}")

    radius = max(1, int(diameter // 2))
    center_scale = float(np.clip(center_scale, 0.0, 1.0))
    edge_scale = float(np.clip(edge_scale, 0.0, 1.0))

    h, w = images.shape[1], images.shape[2]
    cy, cx = h // 2, w // 2

    yy, xx = np.ogrid[:h, :w]
    dist = np.sqrt((yy - cy) ** 2 + (xx - cx) ** 2).astype(np.float32)

    scale = np.ones((h, w), dtype=np.float32)
    inside = dist <= radius
    t = np.zeros((h, w), dtype=np.float32)
    t[inside] = dist[inside] / float(radius)
    scale[inside] = center_scale + t[inside] * (edge_scale - center_scale)

    masked = images.astype(np.float32) * scale[np.newaxis, :, :]

    if np.issubdtype(images.dtype, np.integer):
        info = np.iinfo(images.dtype)
        masked = np.clip(np.rint(masked), info.min, info.max).astype(images.dtype)
    else:
        masked = masked.astype(images.dtype)
    return masked