Source code for ultrasound_metrics.metrics.cnr

"""
Contrast-to-Noise Ratio (CNR) metric for ultrasound image quality assessment.
"""

from typing import Callable

from array_api_compat import array_namespace
from array_api_compat import size as array_size
from beartype import beartype as typechecker
from jaxtyping import jaxtyped

from ultrasound_metrics._utils.array_api import ArrayAPIObj




def get_agg_func(xp: ArrayAPIObj, name: str) -> Callable[[ArrayAPIObj], ArrayAPIObj]:
    """
    Get backend-agnostic aggregation function.

    Parameters
    ----------
    xp
        Array API namespace.
    name
        Name of aggregation function ('MEAN' or 'MEDIAN').

    Returns
    -------
    callable
        Aggregation function.
    """
    if name == "MEAN":
        return lambda arr: xp.mean(arr)
    elif name == "MEDIAN":
        return lambda arr: xp.median(arr)
    else:
        raise ValueError(f"Unknown aggregation function: {name}")



@jaxtyped(typechecker=typechecker)


def compute_cnr(
    values_signal: ArrayAPIObj,
    values_noise: ArrayAPIObj,
    fun_signal: str = "MEAN",
    fun_noise: str = "MEAN",
    *,
    use_signal_variance: bool = True,
    use_noise_variance: bool = True,
) -> float:
    """
    Compute the contrast-to-noise ratio (CNR) between two regions.

    The CNR is commonly used in medical ultrasound to assess image quality [1]_.
    The CNR is computed as [2]_:

    .. math::

        CNR = \\frac{|\\mu_i - \\mu_o|}{\\sqrt{\\sigma_i^2 + \\sigma_o^2}}

    where:

    .. math::

        \\mu_i = E\\{|s_i|^2\\}

        \\mu_o = E\\{|s_o|^2\\}

        \\sigma_i^2 = E\\{(|s_i|^2 - \\mu_i)^2\\}

        \\sigma_o^2 = E\\{(|s_o|^2 - \\mu_o)^2\\}

    with :math:`s_i` and :math:`s_o` representing the signal values inside and outside
    the region of interest, respectively.

    Parameters
    ----------
    values_signal
        Pixel values from the signal region (e.g., inside a lesion).
    values_noise
        Pixel values from the noise/background region (e.g., outside a lesion).
    fun_signal
        Aggregation function for the signal region ("MEAN" or "MEDIAN").
    fun_noise
        Aggregation function for the noise region ("MEAN" or "MEDIAN").
    use_signal_variance
        Whether to include signal variance in the denominator.
    use_noise_variance
        Whether to include noise variance in the denominator.

    Returns
    -------
    float
        The contrast-to-noise ratio.

    References
    ----------
    .. [1] Patterson, M. S., & Foster, F. S. (1983). The improvement and quantitative
           assessment of B-mode images produced by an annular array/cone hybrid.
           Ultrasonic Imaging, 5(3), 195-213.
    .. [2] A. Rodriguez-Molares, O. M. Hoel Rindal, J. D'hooge, S. -E. Måsøy, A. Austeng
           and H. Torp, "The Generalized Contrast-to-Noise Ratio," 2018 IEEE International
           Ultrasonics Symposium (IUS), Kobe, Japan, 2018, pp. 1-4,
           doi: 10.1109/ULTSYM.2018.8580101.
    """
    xp = array_namespace(values_signal)

    if array_size(values_signal) == 0 or array_size(values_noise) == 0:
        raise ValueError("Input arrays must not be empty.")

    # Aggregation functions
    agg_signal = get_agg_func(xp, fun_signal)
    agg_noise = get_agg_func(xp, fun_noise)

    # Match the original: use squared values for mean and variance
    values_signal_sq = xp.abs(values_signal) ** 2
    values_noise_sq = xp.abs(values_noise) ** 2

    mu_signal = agg_signal(values_signal_sq)
    mu_noise = agg_noise(values_noise_sq)
    out = xp.abs(mu_signal - mu_noise)

    if use_noise_variance or use_signal_variance:
        sigma_signal = agg_signal((values_signal_sq - mu_signal) ** 2)
        sigma_noise = agg_noise((values_noise_sq - mu_noise) ** 2)
        denom = xp.sqrt(sigma_signal * int(use_signal_variance) + sigma_noise * int(use_noise_variance))
        out = out / denom

    return float(out)