Merge branch 'refs/heads/main' into fix-voiceprint-dialogue-bug

main/xiaozhi-server/core/connection.py

@@ -11,6 +11,8 @@ import threading
 import traceback
 import subprocess
 import websockets
+import opuslib_next
+import numpy as np
 
 from core.utils.util import (
     extract_json_from_string,
@@ -114,6 +116,7 @@ class ConnectionHandler:
         self.client_abort = False
         self.client_is_speaking = False
         self.client_listen_mode = "auto"
+        self.client_aec = False  # 是否启用了服务端AEC
 
         # 线程任务相关
         self.loop = None  # 在 handle_connection 中获取运行中的事件循环
@@ -153,7 +156,7 @@ class ConnectionHandler:
         # asr相关变量
         # 因为实际部署时可能会用到公共的本地ASR，不能把变量暴露给公共ASR
         # 所以涉及到ASR的变量，需要在这里定义，属于connection的私有变量
-        self.asr_audio = []
+        self.asr_audio = []  # 存储PCM帧列表，供VAD和ASR共享
         self.asr_audio_queue = queue.Queue()
         self.current_speaker = None  # 存储当前说话人
         self.introduced_speakers = set()  # 已"首次引入"的说话人，控制只在首轮带名字
@@ -233,6 +236,9 @@ class ConnectionHandler:
             # 启动超时检查任务
             self.timeout_task = asyncio.create_task(self._check_timeout())
 
+            # 启动AEC缓存清理任务
+            self._aec_cache_cleanup_task = asyncio.create_task(self._check_aec_cache_expiry())
+
             self.welcome_msg = self.config["xiaozhi"]
             self.welcome_msg["session_id"] = self.session_id
 
@@ -368,12 +374,14 @@ class ConnectionHandler:
                 if handled:
                     return
 
-            # 不需要头部处理或没有头部时，直接处理原始消息
-            self.asr_audio_queue.put(message)
+            # 入口处直接解码PCM，避免VAD和ASR重复解码
+            pcm_frame = self._decode_opus_packet(message)
+            if pcm_frame:
+                self.asr_audio_queue.put(pcm_frame)
 
     async def _process_mqtt_audio_message(self, message):
         """
-        处理来自MQTT网关的音频消息，解析16字节头部并提取音频数据
+        处理来自MQTT网关的音频消息，解析16字节头部并提取音频数据，在入队前进行AEC处理
 
         Args:
             message: 包含头部的音频消息
@@ -382,58 +390,192 @@ class ConnectionHandler:
             bool: 是否成功处理了消息
         """
         try:
-            # 提取头部信息
+            # 解析timestamp
             timestamp = int.from_bytes(message[8:12], "big")
-            audio_length = int.from_bytes(message[12:16], "big")
 
-            # 提取音频数据
-            if audio_length > 0 and len(message) >= 16 + audio_length:
-                # 有指定长度，提取精确的音频数据
-                audio_data = message[16 : 16 + audio_length]
-                # 基于时间戳进行排序处理
-                self._process_websocket_audio(audio_data, timestamp)
-                return True
-            elif len(message) > 16:
-                # 没有指定长度或长度无效，去掉头部后处理剩余数据
-                audio_data = message[16:]
-                self.asr_audio_queue.put(audio_data)
+            audio_data = message[16:]
+            # 入口直接解码PCM
+            pcm_frame = self._decode_opus_packet(audio_data)
+            if not pcm_frame:
                 return True
+
+            # AEC处理：如果timestamp>0且启用了AEC
+            if timestamp > 0 and self.client_aec:
+                pcm_frame = self._apply_aec(timestamp, pcm_frame)
+
+            self.asr_audio_queue.put(pcm_frame)
+            return True
         except Exception as e:
             self.logger.bind(tag=TAG).error(f"解析WebSocket音频包失败: {e}")
 
         # 处理失败，返回False表示需要继续处理
         return False
 
-    def _process_websocket_audio(self, audio_data, timestamp):
-        """处理WebSocket格式的音频包"""
-        # 初始化时间戳序列管理
-        if not hasattr(self, "audio_timestamp_buffer"):
-            self.audio_timestamp_buffer = {}
-            self.last_processed_timestamp = 0
-            self.max_timestamp_buffer_size = 20
+    def _apply_aec(self, timestamp: int, pcm_frame: bytes) -> bytes:
+        """应用AEC处理 - 综合算法：互相关延迟估计 + Wiener滤波 + 频谱减法"""
+        try:
+            import numpy as np
 
-        # 如果时间戳是递增的，直接处理
-        if timestamp >= self.last_processed_timestamp:
-            self.asr_audio_queue.put(audio_data)
-            self.last_processed_timestamp = timestamp
+            if not pcm_frame or len(pcm_frame) == 0:
+                return pcm_frame
 
-            # 处理缓冲区中的后续包
-            processed_any = True
-            while processed_any:
-                processed_any = False
-                for ts in sorted(self.audio_timestamp_buffer.keys()):
-                    if ts > self.last_processed_timestamp:
-                        buffered_audio = self.audio_timestamp_buffer.pop(ts)
-                        self.asr_audio_queue.put(buffered_audio)
-                        self.last_processed_timestamp = ts
-                        processed_any = True
-                        break
-        else:
-            # 乱序包，暂存
-            if len(self.audio_timestamp_buffer) < self.max_timestamp_buffer_size:
-                self.audio_timestamp_buffer[timestamp] = audio_data
+            if not hasattr(self, "aec_audio_cache") or not self.aec_audio_cache:
+                return pcm_frame
+
+            mic_audio = np.frombuffer(pcm_frame, dtype=np.int16).astype(np.float32)
+            mic_rms = np.sqrt(np.mean(mic_audio ** 2))
+
+            if mic_rms < 100:
+                return pcm_frame
+
+            # 初始化
+            if not hasattr(self, "_aec_filter_state"):
+                self._aec_filter_state = np.zeros(512)  # 滤波器状态
+                self._aec_delay_ms = 200
+                self._aec_last_delay = 200
+                self._aec_coherence_history = []
+
+            sorted_timestamps = sorted(self.aec_audio_cache.keys())
+            if len(sorted_timestamps) < 2:
+                return pcm_frame
+
+            # ========== 匹配参考帧（对数功率谱匹配） ==========
+            n = len(mic_audio)
+            sorted_timestamps = sorted(self.aec_audio_cache.keys())
+
+            # 找最接近的timestamp作为起点
+            closest_idx = min(range(len(sorted_timestamps)), key=lambda i: abs(sorted_timestamps[i] - timestamp))
+
+            # 对数功率谱相关性计算
+            def log_powerSpectrum_corr(audio1, audio2):
+                window = np.hanning(len(audio1))
+                fft1 = np.fft.rfft(audio1 * window)
+                fft2 = np.fft.rfft(audio2 * window)
+                psd1 = np.abs(fft1) ** 2
+                psd2 = np.abs(fft2) ** 2
+                log1 = 10 * np.log10(psd1 + 1e-8)
+                log2 = 10 * np.log10(psd2 + 1e-8)
+                P_xy = np.dot(log1, log2)
+                P_xx = np.dot(log1, log1)
+                P_yy = np.dot(log2, log2)
+                return abs(P_xy) / (np.sqrt(P_xx) * np.sqrt(P_yy) + 1e-8)
+
+            # 用对数功率谱匹配找最佳帧：前后各找2帧
+            best_corr = -1
+            best_ref_idx = closest_idx
+
+            for offset in range(-2, 3):  # T-2, T-1, T, T+1, T+2
+                test_idx = closest_idx + offset
+                if test_idx < 0 or test_idx >= len(sorted_timestamps):
+                    continue
+                test_ts = sorted_timestamps[test_idx]
+                test_ref = np.frombuffer(self.aec_audio_cache[test_ts], dtype=np.int16).astype(np.float32)
+                test_ref_rms = np.sqrt(np.mean(test_ref ** 2))
+                if test_ref_rms < 50:
+                    continue
+
+                # 对数功率谱相关性
+                corr = log_powerSpectrum_corr(mic_audio, test_ref)
+
+                if corr > best_corr:
+                    best_corr = corr
+                    best_ref_idx = test_idx
+
+            best_ts = sorted_timestamps[best_ref_idx]
+
+            best_ref = np.frombuffer(self.aec_audio_cache[best_ts], dtype=np.int16).astype(np.float32)
+            ref_rms = np.sqrt(np.mean(best_ref ** 2))
+
+            if ref_rms < 50:
+                return pcm_frame
+
+            # 对齐参考信号（直接截取相同长度）
+            aligned_ref = best_ref[:n]
+            if len(aligned_ref) < n:
+                aligned_ref = np.pad(aligned_ref, (0, n - len(aligned_ref)))
+
+            # ========== 频域 AEC 处理（谱减法） ==========
+            # 时域信号经过声学路径后相位失真，导致时域相关性低且P_xy正负不定
+            # 频域幅度谱不受相位影响，对数功率谱相关性稳定在0.97+
+            # 公式：result_mag = max(|mic_fft| - |ref_fft| * scale * coef, 0)
+
+            window = np.hanning(n)
+            mic_fft = np.fft.rfft(mic_audio * window)
+            ref_fft = np.fft.rfft(aligned_ref * window)
+
+            mic_mag = np.abs(mic_fft)
+            ref_mag = np.abs(ref_fft)
+            mic_phase = np.angle(mic_fft)
+
+            # 频域计算回声比例 scale
+            scale = np.sum(mic_mag * ref_mag) / (np.dot(ref_mag, ref_mag) + 1e-8)
+
+            # 使用对数功率谱相关性作为coh
+            coh = best_corr
+            ref_power = ref_rms
+
+            # 自适应系数：根据scale和coh动态调整
+            # scale大（回声强）-> coef大；coh高（匹配准）-> coef大
+            raw_coef = 1.0 + scale * 3 + (coh - 0.97) * 30
+            coef = max(0.5, min(3.0, raw_coef))
+
+            # 谱减法（过减 + 半波整流）
+            if ref_power < 50:
+                output = mic_audio
+                coef = 0
+                gain = 0
             else:
-                self.asr_audio_queue.put(audio_data)
+                echo_mag = ref_mag * scale * coef
+                # 过减：使用较大系数抑制回声
+                over_subtract = 1.5
+                result_mag = np.maximum(mic_mag - echo_mag * over_subtract, mic_mag * 0.1)
+
+                # 保留相位重建信号
+                result_fft = result_mag * np.exp(1j * mic_phase)
+                output = np.fft.irfft(result_fft, n)
+                gain = scale * coef
+
+                # 高置信度是纯回声时，再压一下确保VAD检测不到
+                if coh >= 0.97 and ref_power > 500:
+                    output = output * 0.3
+
+            # 后处理：限幅
+            output = np.clip(output, -32768, 32767)
+
+            # 转换为bytes
+            result = output.astype(np.int16).tobytes()
+
+            return result
+
+        except Exception as e:
+            self.logger.bind(tag=TAG).warning(f"[AEC] 处理失败: {e}")
+            return pcm_frame
+
+    def _decode_opus_packet(self, opus_packet: bytes) -> bytes:
+        """
+        解码Opus数据包为PCM数据
+
+        Args:
+            opus_packet: Opus编码的音频数据
+
+        Returns:
+            bytes: 解码后的PCM数据，失败返回None
+        """
+        try:
+            if not opus_packet or len(opus_packet) == 0:
+                return None
+
+            self._init_connection_state(self)
+            pcm_frame = self._connection_opus_decoder.decode(opus_packet, 960)
+            return pcm_frame
+        except Exception as e:
+            self.logger.bind(tag=TAG).debug(f"Opus解码失败: {e}")
+            return None
+
+    def _init_connection_state(self, conn):
+        """为连接初始化独立的Opus解码器"""
+        if not hasattr(conn, "_connection_opus_decoder"):
+            conn._connection_opus_decoder = opuslib_next.Decoder(16000, 1)
 
     async def handle_restart(self, message):
         """处理服务器重启请求"""
@@ -1413,6 +1555,13 @@ class ConnectionHandler:
             ):
                 self.vad.release_conn_resources(self)
 
+            # 清理opus解码器
+            if hasattr(self, "_connection_opus_decoder"):
+                try:
+                    delattr(self, "_connection_opus_decoder")
+                except Exception:
+                    pass
+
             # 清理音频缓冲区
             if hasattr(self, "audio_buffer"):
                 self.audio_buffer.clear()
@@ -1426,6 +1575,20 @@ class ConnectionHandler:
                     pass
                 self.timeout_task = None
 
+            # 取消AEC缓存清理任务
+            if hasattr(self, "_aec_cache_cleanup_task") and self._aec_cache_cleanup_task and not self._aec_cache_cleanup_task.done():
+                self._aec_cache_cleanup_task.cancel()
+                try:
+                    await self._aec_cache_cleanup_task
+                except asyncio.CancelledError:
+                    pass
+                self._aec_cache_cleanup_task = None
+
+            # 清理AEC缓存
+            if hasattr(self, "aec_audio_cache"):
+                self.aec_audio_cache.clear()
+                self.aec_audio_cache_time.clear()
+
             # 清理工具处理器资源
             if hasattr(self, "func_handler") and self.func_handler:
                 try:
@@ -1589,6 +1752,26 @@ class ConnectionHandler:
         finally:
             self.logger.bind(tag=TAG).info("超时检查任务已退出")
 
+    async def _check_aec_cache_expiry(self):
+        """定期清理过期的AEC缓存"""
+        try:
+            while not self.stop_event.is_set():
+                if hasattr(self, "aec_audio_cache") and self.aec_audio_cache:
+                    current_time = time.time()
+                    expired_keys = [
+                        ts for ts, cache_time in list(self.aec_audio_cache_time.items())
+                        if current_time - cache_time > 120  # 2分钟过期
+                    ]
+                    for ts in expired_keys:
+                        self.aec_audio_cache.pop(ts, None)
+                        self.aec_audio_cache_time.pop(ts, None)
+                    if expired_keys:
+                        self.logger.bind(tag=TAG).debug(f"[AEC] 清理过期缓存 {len(expired_keys)} 条")
+                # 每30秒检查一次
+                await asyncio.sleep(30)
+        except Exception as e:
+            self.logger.bind(tag=TAG).error(f"AEC缓存清理任务出错: {e}")
+
     @staticmethod
     def _extract_direct_answer_response(arguments_str):
         """从 direct_answer 的参数中提取 response 值。