processing_streaming_mel.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
流式Mel特征处理器

用于实时音频流的Mel频谱特征提取，支持chunk-based处理。
支持配置CNN冗余以保证与离线处理的一致性。
"""

import logging
from typing import Dict
from typing import Optional
from typing import Tuple

import numpy as np
import torch

from .processing_audio_minicpma import MiniCPMAAudioProcessor

logger = logging.getLogger(__name__)


class StreamingMelProcessorExact:
    """
    严格离线等价的流式Mel处理器。

    思路：
    - 累积全部历史音频到缓冲；每次新增后用同一个 feature_extractor 计算整段 mel。
    - 只输出"已稳定"的帧：帧中心不依赖未来（右侧）上下文，即 center + n_fft//2 <= 当前缓冲长度。
    - 结束时（flush）再输出最后一批帧，确保与离线全量计算完全一致。

    代价：每次会对累积缓冲做一次特征提取（可按需优化为增量）。
    """

    def __init__(
        self,
        feature_extractor: MiniCPMAAudioProcessor,
        chunk_ms: int = 100,
        first_chunk_ms: Optional[int] = None,
        sample_rate: int = 16000,
        n_fft: int = 400,
        hop_length: int = 160,
        n_mels: int = 80,
        verbose: bool = False,
        cnn_redundancy_ms: int = 10,  # （以ms给定，通常10ms=1帧）
        # --- 滑窗参数（Trigger模式） ---
        enable_sliding_window: bool = False,  # 是否启用滑窗
        slide_trigger_seconds: float = 30.0,  # 触发滑窗的缓冲区秒数阈值
        slide_stride_seconds: float = 10.0,  # 每次滑窗移动的秒数
    ):
        self.feature_extractor = feature_extractor
        self.chunk_ms = chunk_ms
        self.first_chunk_ms = first_chunk_ms if first_chunk_ms is not None else chunk_ms
        self.sample_rate = sample_rate
        self.n_fft = n_fft
        self.hop_length = hop_length
        self.n_mels = n_mels
        self.verbose = verbose

        self.chunk_samples = int(round(chunk_ms * sample_rate / 1000))
        self.chunk_frames = self.chunk_samples // hop_length
        # 对齐到 hop_length 的整数倍，避免帧边界不齐
        hop = self.hop_length
        raw_first_samples = int(round(self.first_chunk_ms * sample_rate / 1000))
        aligned_first = max(hop, (raw_first_samples // hop) * hop)
        self.first_chunk_samples = aligned_first
        self.half_window = n_fft // 2  # 需要的右侧上下文

        # 冗余帧数（以帧为单位），<=1帧：10ms → 1帧
        self.cnn_redundancy_ms = cnn_redundancy_ms
        self.cnn_redundancy_samples = int(cnn_redundancy_ms * sample_rate / 1000)
        self.cnn_redundancy_frames = max(0, self.cnn_redundancy_samples // hop_length)

        # --- 滑窗配置（Trigger模式） ---
        self.enable_sliding_window = enable_sliding_window
        self.trigger_seconds = slide_trigger_seconds
        self.slide_seconds = slide_stride_seconds

        # --- 位移/基准（全局帧坐标） ---
        self.left_samples_dropped = 0  # 已从左侧丢弃的样本数
        self.base_T = 0  # 当前 mel_full[:, :, 0] 对应的"全局帧"下标

        self.reset()

    def reset(self):
        self.buffer = np.zeros(0, dtype=np.float32)
        self.last_emitted_T = 0
        self.total_samples_processed = 0
        self.chunk_count = 0
        self.is_first = True
        self.left_samples_dropped = 0
        self.base_T = 0

    def get_chunk_size(self) -> int:
        return self.first_chunk_samples if self.is_first else self.chunk_samples

    def get_expected_output_frames(self) -> int:
        raise NotImplementedError("get_expected_output_frames is not implemented")

    def _extract_full(self) -> torch.Tensor:
        # 当缓冲长度小于 n_fft 时，Whisper 的内部 STFT 在 center=True 且 pad 模式下会报错
        # （pad 大于输入长度）。此时本来也没有稳定帧可输出，所以直接返回空特征。
        if len(self.buffer) < self.n_fft:
            raise ValueError(f"buffer length is shorter than n_fft {len(self.buffer)} < {self.n_fft}")
        # 如果 buffer 长度 小于 5s 的话，用 set_spac_log_norm(log_floor_db=-10) 或者 上一次缓存的结果
        if len(self.buffer) < 5 * self.sample_rate:
            # TODO: 这里最好的还是 做一些 实验选择 一个 最好的，现在这个 是通过 经验 选择的， 可以看 MiniCPMAAudioProcessor 的 main 实现
            self.feature_extractor.set_spac_log_norm(log_floor_db=-10)
        # 如果 buffer 长度 大于 5s 的话，用 set_spac_log_norm(dynamic_range_db=8)
        else:
            self.feature_extractor.set_spac_log_norm(dynamic_range_db=8)
        feats = self.feature_extractor(
            self.buffer,
            sampling_rate=self.sample_rate,
            return_tensors="pt",
            padding=False,
        )
        return feats.input_features  # [1, 80, T]

    def _stable_frames_count(self) -> int:
        # 已稳定帧数 = floor((len(buffer) - half_window) / hop) + 1，最小为0
        L = int(self.buffer.shape[0])
        if L <= 0:
            return 0
        if L < self.half_window:
            return 0
        return max(0, (L - self.half_window) // self.hop_length + 1)

    def _maybe_slide_buffer(self):
        """Trigger模式滑窗：当缓冲区达到触发阈值时，滑动固定长度的窗口。"""
        if not self.enable_sliding_window:
            return

        sr = self.sample_rate
        hop = self.hop_length
        L = len(self.buffer)

        # 将秒数转换为样本数
        trigger_samples = int(self.trigger_seconds * sr)
        stride_samples = int(self.slide_seconds * sr)

        # 检查是否达到触发阈值
        if L < trigger_samples:
            return

        # 计算需要丢弃的样本数（固定滑动 stride_samples）
        drop = stride_samples

        # 不能丢掉后续发射还需要的左侧上下文
        # 在trigger模式下，我们只需要保护最小必要的数据
        # 即：确保不丢弃未来可能需要的帧
        last_emitted_local = self.last_emitted_T - self.base_T

        # 只保护必要的上下文（例如，最近的1秒数据）
        min_keep_seconds = 1.0  # 保留至少1秒的数据以确保处理的连续性
        min_keep_samples = int(min_keep_seconds * sr)

        # guard_samples 是我们必须保留的最小样本数
        guard_samples = min(min_keep_samples, L - drop)

        # 限制：不得越过安全边界；并对齐 hop
        max_allowed_drop = max(0, L - guard_samples)
        drop = min(drop, max_allowed_drop)
        drop = (drop // hop) * hop

        if drop <= 0:
            return

        # 真正丢弃 & 更新基准
        self.buffer = self.buffer[drop:]
        self.left_samples_dropped += drop
        self.base_T += drop // hop

        if self.verbose:
            print(
                f"[Slide] Trigger模式: drop={drop/sr:.2f}s samples, base_T={self.base_T}, buffer_after={len(self.buffer)/sr:.2f}s"
            )

    def process(self, audio_chunk: np.ndarray, is_last_chunk: bool = False) -> Tuple[torch.Tensor, Dict]:
        self.chunk_count += 1
        # 追加到缓冲
        if len(self.buffer) == 0:
            self.buffer = audio_chunk.astype(np.float32, copy=True)
        else:
            self.buffer = np.concatenate([self.buffer, audio_chunk.astype(np.float32, copy=True)])

        # --- 滑窗处理 ---
        self._maybe_slide_buffer()

        # 全量提取（针对当前窗口）
        mel_full = self._extract_full()
        T_full = mel_full.shape[-1]  # 当前窗口的局部帧数
        stable_T = min(T_full, self._stable_frames_count())  # 局部可稳定帧
        stable_T_global = self.base_T + stable_T  # 映射到全局帧坐标

        # 计划本次发射的核心帧（全局坐标）
        core_start_g = self.last_emitted_T
        core_end_g = core_start_g + self.chunk_frames
        required_stable_g = core_end_g + self.cnn_redundancy_frames

        if self.verbose:
            print(
                f"[Exact] buffer_len={len(self.buffer)} samples, T_full(local)={T_full}, "
                f"stable_T(local)={stable_T}, base_T={self.base_T}, "
                f"stable_T(global)={stable_T_global}, last_emitted={self.last_emitted_T}"
            )

        if stable_T_global >= required_stable_g or is_last_chunk:
            emit_start_g = max(0, core_start_g - self.cnn_redundancy_frames)
            emit_end_g = core_end_g + self.cnn_redundancy_frames

            # 全局 -> 局部索引
            emit_start = max(0, emit_start_g - self.base_T)
            emit_end = emit_end_g - self.base_T
            emit_start = max(0, min(emit_start, T_full))
            emit_end = max(emit_start, min(emit_end, T_full))

            mel_output = mel_full[:, :, emit_start:emit_end]
            self.last_emitted_T = core_end_g  # 仅推进核心帧指针（全局）
        else:
            mel_output = mel_full[:, :, 0:0]

        self.total_samples_processed += len(audio_chunk)
        self.is_first = False

        info = {
            "type": "exact_chunk",
            "chunk_number": self.chunk_count,
            "emitted_frames": mel_output.shape[-1],
            "stable_T": stable_T,
            "T_full": T_full,
            "base_T": self.base_T,
            "stable_T_global": stable_T_global,
            "buffer_len_samples": int(self.buffer.shape[0]),
            "left_samples_dropped": self.left_samples_dropped,
            "core_start": core_start_g,  # 如果保留原字段名，这里用全局值
            "core_end": core_end_g,  # 同上
        }
        return mel_output, info

    def flush(self) -> torch.Tensor:
        """在流结束时调用，输出剩余未发出的帧，保证与离线一致（按全局坐标计算）。"""
        if len(self.buffer) == 0:
            return torch.zeros(1, 80, 0)

        mel_full = self._extract_full()
        T_local = mel_full.shape[-1]
        T_global = self.base_T + T_local

        if self.last_emitted_T < T_global:
            start_l = max(0, self.last_emitted_T - self.base_T)
            tail = mel_full[:, :, start_l:]
            self.last_emitted_T = T_global
            if self.verbose:
                print(f"[Exact] flush {tail.shape[-1]} frames (T_global={T_global})")
            return tail
        return mel_full[:, :, 0:0]

    def get_config(self) -> Dict:
        return {
            "chunk_ms": self.chunk_ms,
            "first_chunk_ms": self.first_chunk_ms,
            "effective_first_chunk_ms": self.first_chunk_samples / self.sample_rate * 1000.0,
            "sample_rate": self.sample_rate,
            "n_fft": self.n_fft,
            "hop_length": self.hop_length,
            "cnn_redundancy_ms": self.cnn_redundancy_ms,
            "cnn_redundancy_frames": self.cnn_redundancy_frames,
            "enable_sliding_window": self.enable_sliding_window,
            "trigger_seconds": self.trigger_seconds,
            "slide_seconds": self.slide_seconds,
        }

    def get_state(self) -> Dict:
        return {
            "chunk_count": self.chunk_count,
            "last_emitted_T": self.last_emitted_T,
            "total_samples_processed": self.total_samples_processed,
            "buffer_len": int(self.buffer.shape[0]),
            "base_T": self.base_T,
            "left_samples_dropped": self.left_samples_dropped,
        }

    def get_snapshot(self) -> Dict:
        """获取完整状态快照（包括 buffer），用于抢跑恢复

        Returns:
            包含完整状态的字典，可用于 restore_snapshot 恢复
        """
        buffer_copy = self.buffer.copy()
        snapshot = {
            "chunk_count": self.chunk_count,
            "last_emitted_T": self.last_emitted_T,
            "total_samples_processed": self.total_samples_processed,
            "buffer": buffer_copy,
            "base_T": self.base_T,
            "left_samples_dropped": self.left_samples_dropped,
            "is_first": self.is_first,
            # 保存 feature_extractor 的状态（关键：确保 mel 特征提取的确定性）
            "fe_dynamic_log_norm": getattr(self.feature_extractor, "dynamic_log_norm", None),
            "fe_dynamic_range_db": getattr(self.feature_extractor, "dynamic_range_db", None),
            "fe_log_floor_db": getattr(self.feature_extractor, "log_floor_db", None),
        }
        logger.debug(
            "[MelProcessor] Created snapshot: chunk_count=%d, last_emitted_T=%d, "
            "buffer_len=%d, buffer_sum=%.6f, total_samples=%d",
            self.chunk_count,
            self.last_emitted_T,
            len(buffer_copy),
            float(buffer_copy.sum()) if len(buffer_copy) > 0 else 0.0,
            self.total_samples_processed,
        )
        return snapshot

    def restore_snapshot(self, snapshot: Dict) -> None:
        """从快照恢复状态

        Args:
            snapshot: 由 get_snapshot 返回的快照字典
        """
        # 记录恢复前的状态
        prev_state = {
            "chunk_count": self.chunk_count,
            "last_emitted_T": self.last_emitted_T,
            "buffer_len": len(self.buffer),
        }

        # 恢复状态
        self.chunk_count = snapshot["chunk_count"]
        self.last_emitted_T = snapshot["last_emitted_T"]
        self.total_samples_processed = snapshot["total_samples_processed"]
        self.buffer = snapshot["buffer"].copy()  # 复制 buffer
        self.base_T = snapshot["base_T"]
        self.left_samples_dropped = snapshot["left_samples_dropped"]
        self.is_first = snapshot["is_first"]

        # 恢复 feature_extractor 的状态（关键：确保 mel 特征提取的确定性）
        if snapshot.get("fe_dynamic_log_norm") is not None:
            self.feature_extractor.dynamic_log_norm = snapshot["fe_dynamic_log_norm"]
        if snapshot.get("fe_dynamic_range_db") is not None:
            self.feature_extractor.dynamic_range_db = snapshot["fe_dynamic_range_db"]
        if snapshot.get("fe_log_floor_db") is not None:
            self.feature_extractor.log_floor_db = snapshot["fe_log_floor_db"]

        logger.info(
            "[MelProcessor] Restored snapshot: chunk_count %d->%d, last_emitted_T %d->%d, "
            "buffer_len %d->%d, total_samples=%d",
            prev_state["chunk_count"],
            self.chunk_count,
            prev_state["last_emitted_T"],
            self.last_emitted_T,
            prev_state["buffer_len"],
            len(self.buffer),
            self.total_samples_processed,
        )