317 lines
12 KiB
Python
317 lines
12 KiB
Python
import asyncio
|
||
import json
|
||
import os
|
||
import tempfile
|
||
import wave
|
||
import websockets
|
||
import aiohttp
|
||
import logging
|
||
|
||
from dotenv import load_dotenv
|
||
|
||
logging.basicConfig(
|
||
level=logging.INFO,
|
||
format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
|
||
force=True
|
||
)
|
||
logger = logging.getLogger("app")
|
||
|
||
|
||
# ========================
|
||
# 配置
|
||
# ========================
|
||
load_dotenv()
|
||
logger.info("正在加载环境变量 ")
|
||
TARGET_WS_URL = os.getenv("TARGET_WS_URL", "ws://localhost:59805/ws/asr")
|
||
LOCAL_HOST = "0.0.0.0"
|
||
LOCAL_PORT = 10095
|
||
ASR_LOCAL_URL = os.getenv("ASR_LOCAL_URL", "http://172.18.127.124:6688/asr")
|
||
|
||
SAMPLE_RATE = 16000
|
||
SAMPLE_WIDTH = 2 # 16-bit = 2 bytes
|
||
CHANNELS = 1
|
||
|
||
# 计算每秒的字节数
|
||
BYTES_PER_SECOND = SAMPLE_RATE * SAMPLE_WIDTH * CHANNELS
|
||
# 设定最小触发时长（秒）
|
||
MIN_DURATION_THRESHOLD = 3.0
|
||
|
||
config = {
|
||
"chunk_size": [10, 10, 10],
|
||
"wav_name": "h5",
|
||
"is_speaking": True,
|
||
"wav_format": "pcm",
|
||
"chunk_interval": 10,
|
||
"itn": True,
|
||
"mode": "2pass",
|
||
"hotwords": "",
|
||
}
|
||
|
||
|
||
# ========================
|
||
# 新增：声纹会话缓存管理类 (带超时机制 + 日志)
|
||
# ========================
|
||
class SpeakerSession:
|
||
def __init__(self):
|
||
self.last_speaker = None
|
||
self.buffer = []
|
||
self.accumulated_duration = 0.0
|
||
self.lock = asyncio.Lock()
|
||
self.MAX_BUFFER_DURATION = 3.0
|
||
|
||
async def process_and_send(self, new_msg, websocket, duration=0.0, force_flush=False):
|
||
"""
|
||
duration: 当前这段消息对应的音频时长(秒)
|
||
"""
|
||
async with self.lock:
|
||
# === 1. 强制清空 (连接断开时) ===
|
||
if force_flush:
|
||
if self.buffer:
|
||
logger.info(f"🧹 [最终清空] 准备推送 {len(self.buffer)} 条缓存...")
|
||
for msg in self.buffer:
|
||
text_content = msg.get("full_text", "")
|
||
logger.info(f" └── 🚀 发送: {text_content}") # 打印发送内容
|
||
await websocket.send(json.dumps(msg, ensure_ascii=False))
|
||
self.buffer = []
|
||
self.accumulated_duration = 0.0
|
||
return
|
||
|
||
# 获取当前消息的说话人
|
||
segments = new_msg.get("segments", [])
|
||
current_speaker = segments[0].get("spk") if segments else "unknown"
|
||
|
||
# === 2. 如果是同一个人（或者刚开始） ===
|
||
if self.last_speaker is None or current_speaker == self.last_speaker:
|
||
self.buffer.append(new_msg)
|
||
self.accumulated_duration += duration
|
||
self.last_speaker = current_speaker
|
||
|
||
# 🌟 检查是否超时 20s
|
||
if self.accumulated_duration >= self.MAX_BUFFER_DURATION:
|
||
logger.info(f"⏰ [超时推送] 同人 ({current_speaker}) 积攒 {self.accumulated_duration:.2f}s >= 8s")
|
||
|
||
for msg in self.buffer:
|
||
text_content = msg.get("full_text", "")
|
||
logger.info(f" └── 🚀 发送: {text_content}") # 打印发送内容
|
||
await websocket.send(json.dumps(msg, ensure_ascii=False))
|
||
|
||
# 清空缓存，重置时长
|
||
self.buffer = []
|
||
self.accumulated_duration = 0.0
|
||
else:
|
||
# 仅缓存，不发送
|
||
text_preview = new_msg.get("full_text", "")
|
||
logger.info(f"📥 [缓存中] Spk:{current_speaker} | 文本:{text_preview} | 累计:{self.accumulated_duration:.2f}s")
|
||
|
||
# === 3. 如果换人了 ===
|
||
else:
|
||
logger.info(f"🔄 [换人推送] {self.last_speaker} -> {current_speaker}")
|
||
|
||
# A. 推送上一位说话人的所有积攒消息
|
||
for msg in self.buffer:
|
||
text_content = msg.get("full_text", "")
|
||
logger.info(f" └── 🚀 发送: {text_content}") # 打印发送内容
|
||
await websocket.send(json.dumps(msg, ensure_ascii=False))
|
||
|
||
# B. 清空并开始缓存新人的消息
|
||
self.buffer = []
|
||
self.buffer.append(new_msg)
|
||
self.accumulated_duration = duration
|
||
self.last_speaker = current_speaker
|
||
|
||
# 打印新人入队日志
|
||
new_text = new_msg.get("full_text", "")
|
||
logger.info(f"📥 [新缓存] Spk:{current_speaker} | 文本:{new_text}")
|
||
|
||
# ========================
|
||
# 调用本地 ASR 接口（带说话人识别）
|
||
# ========================
|
||
async def call_local_asr_with_audio(wav_path: str, session: aiohttp.ClientSession):
|
||
try:
|
||
with open(wav_path, "rb") as f:
|
||
data = aiohttp.FormData()
|
||
data.add_field(
|
||
"audio_file",
|
||
f,
|
||
filename=os.path.basename(wav_path),
|
||
content_type="audio/wav"
|
||
)
|
||
params = {
|
||
"spk_diarization": "true",
|
||
"spk_threshold": "0.6",
|
||
"output": "json"
|
||
}
|
||
async with session.post(ASR_LOCAL_URL, data=data, params=params, timeout=aiohttp.ClientTimeout(total=60)) as resp:
|
||
if resp.status == 200:
|
||
return await resp.json()
|
||
else:
|
||
text = await resp.text()
|
||
logger.error(f"❌ 本地 ASR 接口错误: {resp.status} - {text}")
|
||
return None
|
||
except Exception as e:
|
||
logger.error(f"❌ 调用本地 ASR 异常: {e}")
|
||
return None
|
||
|
||
|
||
# ========================
|
||
# 后台处理：音频 → 本地 ASR → Session 缓存判断
|
||
# ========================
|
||
async def process_audio_with_local_asr(audio_bytes, asr_text, websocket, session, speaker_session):
|
||
if not audio_bytes:
|
||
return
|
||
|
||
wav_path = None
|
||
try:
|
||
# 1. 计算这段音频的时长，用于 Session 判断超时
|
||
audio_duration = len(audio_bytes) / BYTES_PER_SECOND
|
||
|
||
# 2. 保存为 WAV
|
||
with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as tmp:
|
||
wav_path = tmp.name
|
||
|
||
with wave.open(wav_path, 'wb') as wf:
|
||
wf.setnchannels(CHANNELS)
|
||
wf.setsampwidth(SAMPLE_WIDTH)
|
||
wf.setframerate(SAMPLE_RATE)
|
||
wf.writeframes(audio_bytes)
|
||
|
||
# 3. 调用本地 ASR 接口
|
||
asr_result = await call_local_asr_with_audio(wav_path, session)
|
||
if not asr_result or not asr_result.get("success"):
|
||
logger.error("❌ 本地 ASR (带声纹) 处理失败")
|
||
return
|
||
|
||
# 4. 构建消息
|
||
segments = asr_result.get("segments", [])
|
||
full_text = asr_result.get("text", "")
|
||
|
||
msg = {
|
||
"mode": "offline-speaker",
|
||
"full_text": full_text,
|
||
"segments": segments,
|
||
"original_asr_text": asr_text,
|
||
"type": "asr_with_speaker",
|
||
}
|
||
|
||
# 5. 🌟 传入 duration 参数
|
||
await speaker_session.process_and_send(msg, websocket, duration=audio_duration)
|
||
|
||
except websockets.exceptions.ConnectionClosed:
|
||
logger.warning("⚠️ WebSocket 连接已关闭")
|
||
except Exception as e:
|
||
logger.error(f"❌ 后台处理异常: {e}")
|
||
finally:
|
||
if wav_path and os.path.exists(wav_path):
|
||
os.unlink(wav_path)
|
||
|
||
|
||
# ========================
|
||
# 主 WebSocket 转发逻辑
|
||
# ========================
|
||
async def forward_audio(websocket):
|
||
logger.info("🟢 前端已连接。")
|
||
|
||
current_audio_chunks = []
|
||
local_accumulated_audio = bytearray()
|
||
|
||
speaker_session = SpeakerSession()
|
||
|
||
lock = asyncio.Lock()
|
||
|
||
async with aiohttp.ClientSession() as session:
|
||
try:
|
||
async with websockets.connect(TARGET_WS_URL) as target_ws:
|
||
logger.info(f"🔗 连接到远程 ASR: {TARGET_WS_URL}")
|
||
await target_ws.send(json.dumps(config))
|
||
|
||
# ---------------------------
|
||
# 任务1: 前端 -> 远程 ASR
|
||
# ---------------------------
|
||
async def frontend_to_target():
|
||
async for message in websocket:
|
||
if isinstance(message, bytes):
|
||
async with lock:
|
||
current_audio_chunks.append(message)
|
||
await target_ws.send(message)
|
||
else:
|
||
await target_ws.send(message)
|
||
|
||
# ---------------------------
|
||
# 任务2: 远程 ASR -> 前端
|
||
# ---------------------------
|
||
async def target_to_frontend():
|
||
async for message in target_ws:
|
||
# 实时推送
|
||
await websocket.send(message)
|
||
|
||
if isinstance(message, str):
|
||
try:
|
||
data = json.loads(message)
|
||
mode = data.get("mode")
|
||
|
||
if mode == "2pass-offline":
|
||
asr_text = data.get("text", "").strip()
|
||
|
||
async with lock:
|
||
chunk_bytes = b"".join(current_audio_chunks)
|
||
current_audio_chunks.clear()
|
||
|
||
local_accumulated_audio.extend(chunk_bytes)
|
||
current_duration = len(local_accumulated_audio) / BYTES_PER_SECOND
|
||
|
||
if current_duration < MIN_DURATION_THRESHOLD:
|
||
logger.info(f"⏳ 音频拼接中: {current_duration:.2f}s")
|
||
else:
|
||
logger.info(f"⚡ 触发本地声纹识别: {current_duration:.2f}s")
|
||
|
||
audio_to_send = bytes(local_accumulated_audio)
|
||
local_accumulated_audio.clear()
|
||
|
||
asyncio.create_task(
|
||
process_audio_with_local_asr(audio_to_send, asr_text, websocket, session, speaker_session)
|
||
)
|
||
|
||
elif mode == "2pass-online":
|
||
pass
|
||
|
||
except json.JSONDecodeError:
|
||
pass
|
||
|
||
await asyncio.gather(frontend_to_target(), target_to_frontend())
|
||
|
||
except websockets.exceptions.ConnectionClosedOK:
|
||
logger.info("🔗 远程 ASR 连接正常关闭。")
|
||
except Exception as e:
|
||
logger.error(f"❌ 连接异常: {e}")
|
||
finally:
|
||
logger.info("🔴 前端断开连接，清理资源...")
|
||
|
||
if len(local_accumulated_audio) > 0:
|
||
final_duration = len(local_accumulated_audio) / BYTES_PER_SECOND
|
||
logger.info(f"⚠️ 处理断开前的剩余音频: {final_duration:.2f}s")
|
||
try:
|
||
audio_to_send = bytes(local_accumulated_audio)
|
||
local_accumulated_audio.clear()
|
||
await process_audio_with_local_asr(audio_to_send, "Final-Flush", websocket, session, speaker_session)
|
||
except Exception as e:
|
||
logger.error(f"❌ 剩余音频处理异常: {e}")
|
||
|
||
# 强制推送最后残留缓存
|
||
try:
|
||
await speaker_session.process_and_send(None, websocket, force_flush=True)
|
||
except Exception as e:
|
||
logger.error(f"❌ 最终缓存推送失败: {e}")
|
||
|
||
|
||
# ========================
|
||
# 启动服务
|
||
# ========================
|
||
async def main():
|
||
logger.info(f"🚀 启动 WebSocket 中转服务: ws://{LOCAL_HOST}:{LOCAL_PORT}")
|
||
logger.info("📝 逻辑: 声纹识别缓存，同人累积满8s或换人时推送")
|
||
server = await websockets.serve(forward_audio, LOCAL_HOST, LOCAL_PORT)
|
||
await server.wait_closed()
|
||
|
||
|
||
if __name__ == "__main__":
|
||
asyncio.run(main()) |