speech_module/audio_preprocess.py

from __future__ import annotations

import io
from typing import Any

import numpy as np
import soundfile as sf

from config import (
    SPEECH_ENERGY_THRESHOLD_FLOOR,
    SPEECH_ENERGY_THRESHOLD_RATIO,
    SPEECH_MIN_SPEECH_RATIO,
    SPEECH_MIN_SPEECH_SECONDS,
    SPEECH_NOISE_CALIBRATION_SECONDS,
    SPEECH_SPECTRAL_FLATNESS_MAX,
    SPEECH_ZERO_CROSSING_MAX,
)


def extract_useful_speech_audio(
    audio_bytes: bytes,
    *,
    vad_model: Any | None = None,
    target_sr: int = 16000,
    vad_threshold: float = 0.5,
    min_speech_seconds: float = SPEECH_MIN_SPEECH_SECONDS,
    min_speech_ratio: float = SPEECH_MIN_SPEECH_RATIO,
    energy_threshold_floor: float = SPEECH_ENERGY_THRESHOLD_FLOOR,
    energy_threshold_ratio: float = SPEECH_ENERGY_THRESHOLD_RATIO,
    spectral_flatness_max: float = SPEECH_SPECTRAL_FLATNESS_MAX,
    zero_crossing_max: float = SPEECH_ZERO_CROSSING_MAX,
    noise_calibration_seconds: float = SPEECH_NOISE_CALIBRATION_SECONDS,
) -> dict[str, object]:
    """Filter audio to useful speech chunks and reject noise-dominant clips."""
    if not audio_bytes:
        return _unavailable_result("empty_audio", "音频输入为空，无法分析。")

    audio, sample_rate = _bytes_to_array(audio_bytes, target_sr=target_sr)
    total_seconds = float(len(audio)) / float(sample_rate) if len(audio) else 0.0
    if total_seconds < 0.25:
        return _unavailable_result("too_short", "音频时长过短，无法稳定识别语音。")

    speech_audio = np.array([], dtype=np.float32)
    speech_seconds = 0.0
    method = "energy_fallback"

    if vad_model is not None:
        try:
            speech_audio, speech_seconds = _extract_speech_with_vad(
                audio,
                sample_rate,
                vad_model,
                threshold=vad_threshold,
            )
            method = "silero_vad"
        except Exception:
            speech_audio = np.array([], dtype=np.float32)
            speech_seconds = 0.0

    if len(speech_audio) == 0:
        speech_audio, speech_seconds = _extract_speech_with_energy(
            audio,
            sample_rate,
            energy_threshold_floor=energy_threshold_floor,
            energy_threshold_ratio=energy_threshold_ratio,
            spectral_flatness_max=spectral_flatness_max,
            zero_crossing_max=zero_crossing_max,
            noise_calibration_seconds=noise_calibration_seconds,
        )

    speech_ratio = speech_seconds / total_seconds if total_seconds > 0 else 0.0
    if speech_seconds < min_speech_seconds or speech_ratio < min_speech_ratio:
        return {
            **_unavailable_result(
                "noise_only",
                "检测到的有效语音过少，当前录音以杂音/静音为主，请重录并靠近麦克风。",
            ),
            "total_seconds": round(total_seconds, 3),
            "speech_seconds": round(speech_seconds, 3),
            "speech_ratio": round(speech_ratio, 4),
            "method": method,
        }

    if len(speech_audio) == 0:
        speech_audio = audio
        speech_seconds = total_seconds
        speech_ratio = 1.0

    return {
        "available": True,
        "error_code": None,
        "warning": None,
        "message": "ok",
        "audio_bytes": _array_to_wav_bytes(_normalize_audio(speech_audio), sample_rate),
        "total_seconds": round(total_seconds, 3),
        "speech_seconds": round(speech_seconds, 3),
        "speech_ratio": round(speech_ratio, 4),
        "method": method,
    }


def _unavailable_result(error_code: str, message: str) -> dict[str, object]:
    return {
        "available": False,
        "error_code": error_code,
        "warning": message,
        "message": message,
        "audio_bytes": None,
        "total_seconds": 0.0,
        "speech_seconds": 0.0,
        "speech_ratio": 0.0,
        "method": None,
    }


def _extract_speech_with_vad(
    audio: np.ndarray,
    sample_rate: int,
    vad_model: Any,
    *,
    threshold: float,
) -> tuple[np.ndarray, float]:
    import torch
    from silero_vad import collect_chunks, get_speech_timestamps

    tensor = torch.from_numpy(audio.astype(np.float32))
    timestamps = get_speech_timestamps(
        tensor,
        vad_model,
        sampling_rate=sample_rate,
        threshold=float(threshold),
        min_speech_duration_ms=250,
        min_silence_duration_ms=120,
        speech_pad_ms=80,
        return_seconds=False,
    )
    if not timestamps:
        return np.array([], dtype=np.float32), 0.0

    speech_tensor = collect_chunks(timestamps, tensor)
    speech_audio = speech_tensor.detach().cpu().numpy().astype(np.float32)
    speech_samples = sum(
        max(0, int(item.get("end", 0)) - int(item.get("start", 0)))
        for item in timestamps
        if isinstance(item, dict)
    )
    speech_seconds = float(speech_samples) / float(sample_rate)
    return speech_audio, speech_seconds


def _extract_speech_with_energy(
    audio: np.ndarray,
    sample_rate: int,
    *,
    energy_threshold_floor: float,
    energy_threshold_ratio: float,
    spectral_flatness_max: float,
    zero_crossing_max: float,
    noise_calibration_seconds: float,
) -> tuple[np.ndarray, float]:
    frame_len = max(1, int(sample_rate * 0.03))
    hop_len = max(1, int(sample_rate * 0.01))
    if len(audio) < frame_len:
        return np.array([], dtype=np.float32), 0.0

    frames = _frame_audio(audio, frame_len, hop_len)
    rms = np.sqrt(np.mean(frames**2, axis=1))
    max_rms = float(np.max(rms))
    noise_samples = min(len(audio), int(sample_rate * max(0.05, noise_calibration_seconds)))
    noise_rms = float(np.sqrt(np.mean(audio[:noise_samples] ** 2))) if noise_samples > 0 else 0.0
    noise_threshold = noise_rms * 2.5 if max_rms > 0 and noise_rms < max_rms * 0.6 else 0.0
    rms_threshold = max(
        float(energy_threshold_floor),
        max_rms * float(energy_threshold_ratio),
        noise_threshold,
    )

    window = np.hanning(frame_len).astype(np.float32)
    spectrum = np.fft.rfft(frames * window[None, :], axis=1)
    power = (np.abs(spectrum) ** 2) + 1e-12
    flatness = np.exp(np.mean(np.log(power), axis=1)) / np.mean(power, axis=1)
    zero_cross = np.mean(np.diff(np.signbit(frames), axis=1), axis=1)

    voiced_mask = (
        (rms >= rms_threshold)
        & (flatness <= float(spectral_flatness_max))
        & (np.abs(zero_cross) <= float(zero_crossing_max))
    )
    if not np.any(voiced_mask):
        return np.array([], dtype=np.float32), 0.0

    voiced_samples = int(np.sum(voiced_mask)) * hop_len
    speech_seconds = float(voiced_samples) / float(sample_rate)
    speech_audio = _collect_segments(audio, voiced_mask, frame_len, hop_len, sample_rate)
    return speech_audio, speech_seconds


def _frame_audio(audio: np.ndarray, frame_len: int, hop_len: int) -> np.ndarray:
    frame_count = 1 + (len(audio) - frame_len) // hop_len
    starts = np.arange(frame_count) * hop_len
    indices = starts[:, None] + np.arange(frame_len)[None, :]
    return audio[indices]


def _collect_segments(
    audio: np.ndarray,
    voiced_mask: np.ndarray,
    frame_len: int,
    hop_len: int,
    sample_rate: int,
) -> np.ndarray:
    pad = int(sample_rate * 0.05)
    segments: list[np.ndarray] = []
    start: int | None = None

    for index, is_voiced in enumerate(voiced_mask.tolist()):
        if is_voiced and start is None:
            start = index
        elif not is_voiced and start is not None:
            left = max(0, start * hop_len - pad)
            right = min(len(audio), index * hop_len + frame_len + pad)
            if right > left:
                segments.append(audio[left:right])
            start = None

    if start is not None:
        left = max(0, start * hop_len - pad)
        right = len(audio)
        if right > left:
            segments.append(audio[left:right])

    if not segments:
        return np.array([], dtype=np.float32)
    return np.concatenate(segments).astype(np.float32)


def _array_to_wav_bytes(audio: np.ndarray, sample_rate: int) -> bytes:
    buffer = io.BytesIO()
    sf.write(buffer, audio.astype(np.float32), sample_rate, format="WAV")
    return buffer.getvalue()


def _normalize_audio(audio: np.ndarray, target_peak: float = 0.85) -> np.ndarray:
    if len(audio) == 0:
        return audio.astype(np.float32)
    peak = float(np.max(np.abs(audio)))
    if peak <= 1e-6:
        return audio.astype(np.float32)
    scale = min(1.0 / peak, float(target_peak) / peak)
    return np.clip(audio.astype(np.float32) * scale, -1.0, 1.0)


def _bytes_to_array(audio_bytes: bytes, target_sr: int = 16000) -> tuple[np.ndarray, int]:
    try:
        import librosa

        bio = io.BytesIO(audio_bytes)
        bio.seek(0)
        audio, _ = librosa.load(bio, sr=target_sr, mono=True)
        return audio.astype(np.float32), target_sr
    except Exception:
        bio = io.BytesIO(audio_bytes)
        bio.seek(0)
        audio, source_sr = sf.read(bio)
        if audio.ndim > 1:
            audio = audio.mean(axis=1)
        if source_sr != target_sr:
            import scipy.signal

            gcd_val = int(np.gcd(source_sr, target_sr))
            audio = scipy.signal.resample_poly(
                audio,
                up=target_sr // gcd_val,
                down=source_sr // gcd_val,
            )
        return audio.astype(np.float32), target_sr