backend/services/quality_metrics.py

"""Quality metrics for audio processing evaluation."""

import numpy as np
from scipy import signal


def compute_snr(original: np.ndarray, processed: np.ndarray) -> float:
    """
    Compute Signal-to-Noise Ratio in dB.

    Args:
        original: Original audio array
        processed: Processed audio array

    Returns:
        SNR in dB (higher = better)
    """
    # Align lengths
    min_len = min(len(original), len(processed))
    original = original[:min_len]
    processed = processed[:min_len]

    # Compute noise as difference
    noise = original - processed

    # Compute powers
    signal_power = np.mean(original ** 2)
    noise_power = np.mean(noise ** 2)

    # Handle edge case
    if noise_power < 1e-10:
        return float('inf')

    snr = 10 * np.log10(signal_power / noise_power)
    return float(snr)


def compute_spectral_similarity(
    original: np.ndarray,
    processed: np.ndarray,
    sr: int
) -> float:
    """
    Compute cosine similarity of magnitude spectrograms.

    Args:
        original: Original audio array
        processed: Processed audio array
        sr: Sample rate

    Returns:
        Similarity score (0-1, higher = more similar)
    """
    # Align lengths
    min_len = min(len(original), len(processed))
    original = original[:min_len]
    processed = processed[:min_len]

    # Compute spectrograms using STFT
    nperseg = min(2048, min_len // 4)  # Window size
    if nperseg < 16:
        return 1.0  # Too short to compute

    _, _, Zxx_orig = signal.stft(original, fs=sr, nperseg=nperseg)
    _, _, Zxx_proc = signal.stft(processed, fs=sr, nperseg=nperseg)

    # Get magnitude spectrograms
    mag_orig = np.abs(Zxx_orig).flatten()
    mag_proc = np.abs(Zxx_proc).flatten()

    # Compute cosine similarity
    dot_product = np.dot(mag_orig, mag_proc)
    norm_orig = np.linalg.norm(mag_orig)
    norm_proc = np.linalg.norm(mag_proc)

    if norm_orig < 1e-10 or norm_proc < 1e-10:
        return 1.0 if norm_orig == norm_proc else 0.0

    similarity = dot_product / (norm_orig * norm_proc)
    return float(np.clip(similarity, 0, 1))