live-app-translator/src/livesub/audio/segmenter.py

"""발화 구간 분할기.

캡처된 연속 오디오를 "한 문장" 단위로 잘라 음성인식에 넘긴다.
외부 의존성 없이 동작하도록 적응형 노이즈 플로어 기반 에너지 VAD를 쓴다.
webrtcvad 가 설치돼 있으면 더 정확한 판정을 위해 함께 사용한다.
"""

from __future__ import annotations

import time
from dataclasses import dataclass

import numpy as np

from ..constants import FRAME_MS, FRAME_SAMPLES, SAMPLE_RATE

try:  # pragma: no cover - 선택 의존성
    import webrtcvad
except ImportError:  # pragma: no cover
    webrtcvad = None  # type: ignore[assignment]


@dataclass
class Segment:
    """음성인식에 넘길 오디오 한 덩어리."""

    audio: np.ndarray
    started_at: float
    ended_at: float
    is_final: bool = True

    @property
    def duration_s(self) -> float:
        return self.audio.size / SAMPLE_RATE


@dataclass
class SegmenterConfig:
    silence_ms: int = 600
    min_segment_ms: int = 400
    max_segment_ms: int = 12_000
    partial_interval_ms: int = 900
    #: 노이즈 플로어 대비 몇 배 이상이면 발화로 볼지
    speech_ratio: float = 3.0
    #: 절대 무음 판정 하한 (RMS)
    absolute_floor: float = 2e-4
    vad_aggressiveness: int = 2
    #: 발화 시작 앞쪽에 붙일 여유분 (앞 음절 잘림 방지)
    lead_in_ms: int = 240


class Segmenter:
    """프레임을 받아 발화 구간을 뱉는 상태 기계.

    `push()` 는 확정 구간(is_final=True) 또는 중간 결과(is_final=False)를
    돌려주거나, 아직 내보낼 게 없으면 None 을 돌려준다.
    """

    def __init__(self, config: SegmenterConfig | None = None) -> None:
        self.config = config or SegmenterConfig()
        self._vad = None
        if webrtcvad is not None:
            self._vad = webrtcvad.Vad(max(0, min(3, self.config.vad_aggressiveness)))

        self._buffer: list[np.ndarray] = []
        self._lead_in: list[np.ndarray] = []
        self._pending = np.zeros(0, dtype=np.float32)
        self._noise_floor = 1e-3
        self._silence_ms = 0
        self._speech_ms = 0
        self._in_speech = False
        self._started_at = 0.0
        self._last_partial_ms = 0
        self._lead_in_frames = max(1, self.config.lead_in_ms // FRAME_MS)

    # --- 입력 ----------------------------------------------------------
    def push(self, chunk: np.ndarray) -> list[Segment]:
        """임의 길이의 오디오를 넣고 완성된 구간 목록을 받는다."""
        out: list[Segment] = []
        if chunk.size:
            self._pending = (
                chunk.astype(np.float32, copy=False)
                if self._pending.size == 0
                else np.concatenate((self._pending, chunk))
            )
        while self._pending.size >= FRAME_SAMPLES:
            frame = self._pending[:FRAME_SAMPLES]
            self._pending = self._pending[FRAME_SAMPLES:]
            segment = self._push_frame(frame)
            if segment is not None:
                out.append(segment)
        return out

    def flush(self) -> Segment | None:
        """남은 버퍼를 강제로 확정 구간으로 만든다 (캡처 종료 시)."""
        if not self._in_speech:
            return None
        return self._finish()

    def reset(self) -> None:
        """발화 상태만 초기화한다.

        `_pending` 은 건드리지 않는다. 아직 프레임으로 쪼개지 못한 '입력' 이라
        여기서 버리면 한 구간을 확정한 직후의 오디오가 통째로 사라진다.
        """
        self._buffer.clear()
        self._lead_in.clear()
        self._silence_ms = 0
        self._speech_ms = 0
        self._in_speech = False
        self._last_partial_ms = 0

    # --- 내부 ----------------------------------------------------------
    def _push_frame(self, frame: np.ndarray) -> Segment | None:
        speech = self._is_speech(frame)
        cfg = self.config

        if not self._in_speech:
            self._lead_in.append(frame)
            if len(self._lead_in) > self._lead_in_frames:
                self._lead_in.pop(0)
            if speech:
                self._in_speech = True
                self._started_at = time.monotonic()
                self._buffer = list(self._lead_in)
                self._lead_in = []
                self._silence_ms = 0
                self._speech_ms = FRAME_MS
                self._last_partial_ms = 0
            return None

        self._buffer.append(frame)
        self._speech_ms += FRAME_MS
        self._silence_ms = 0 if speech else self._silence_ms + FRAME_MS

        buffered_ms = len(self._buffer) * FRAME_MS
        if self._silence_ms >= cfg.silence_ms:
            if buffered_ms - self._silence_ms < cfg.min_segment_ms:
                self.reset()  # 잡음 한 번 튄 것 — 버린다
                return None
            return self._finish()

        if buffered_ms >= cfg.max_segment_ms:
            return self._finish()

        if cfg.partial_interval_ms > 0:
            since = buffered_ms - self._last_partial_ms
            if since >= cfg.partial_interval_ms and buffered_ms >= cfg.min_segment_ms:
                self._last_partial_ms = buffered_ms
                return Segment(
                    audio=np.concatenate(self._buffer),
                    started_at=self._started_at,
                    ended_at=time.monotonic(),
                    is_final=False,
                )
        return None

    def _finish(self) -> Segment:
        audio = np.concatenate(self._buffer) if self._buffer else np.zeros(0, np.float32)
        segment = Segment(
            audio=audio,
            started_at=self._started_at,
            ended_at=time.monotonic(),
            is_final=True,
        )
        self.reset()
        return segment

    def _is_speech(self, frame: np.ndarray) -> bool:
        rms = float(np.sqrt(np.mean(np.square(frame), dtype=np.float64)))
        cfg = self.config

        if rms < cfg.absolute_floor:
            self._noise_floor = min(self._noise_floor, max(rms, 1e-6))
            return False

        energetic = rms > max(self._noise_floor * cfg.speech_ratio, cfg.absolute_floor)
        if not energetic:
            # 무음 구간에서만 노이즈 플로어를 천천히 따라가게 한다.
            self._noise_floor = 0.95 * self._noise_floor + 0.05 * rms

        if self._vad is None:
            return energetic
        try:
            pcm16 = (np.clip(frame, -1.0, 1.0) * 32767.0).astype(np.int16).tobytes()
            return energetic and self._vad.is_speech(pcm16, SAMPLE_RATE)
        except Exception:  # noqa: BLE001 - webrtcvad 는 프레임 길이에 민감
            return energetic