111 lines
4.5 KiB
Python
111 lines
4.5 KiB
Python
import asyncio
|
|
import websockets
|
|
import soundfile as sf
|
|
import uuid
|
|
|
|
# --- 配置 ---
|
|
HOST = "localhost"
|
|
PORT = 11096
|
|
SESSION_ID = str(uuid.uuid4())
|
|
SENDER_URI = f"ws://{HOST}:{PORT}/ws/asr/{SESSION_ID}?mode=sender"
|
|
RECEIVER_URI = f"ws://{HOST}:{PORT}/ws/asr/{SESSION_ID}?mode=receiver"
|
|
|
|
AUDIO_FILE_PATH = "tests/XT_ZZY_denoise.wav" # 确保此测试文件存在且为 16kHz, 16-bit, 单声道
|
|
CHUNK_DURATION_MS = 100 # 每次发送100ms的音频数据
|
|
CHUNK_SIZE = int(16000 * 2 * CHUNK_DURATION_MS / 1000) # 3200 bytes
|
|
|
|
async def run_receiver():
|
|
"""作为接收者连接，并打印收到的所有消息。"""
|
|
print(f"▶️ [Receiver] 尝试连接到: {RECEIVER_URI}")
|
|
try:
|
|
async with websockets.connect(RECEIVER_URI) as websocket:
|
|
print("✅ [Receiver] 连接成功，等待消息...")
|
|
try:
|
|
while True:
|
|
message = await websocket.recv()
|
|
print(f"🎧 [Receiver] 收到结果: {message}")
|
|
except websockets.exceptions.ConnectionClosed as e:
|
|
print(f"✅ [Receiver] 连接已由服务器正常关闭: {e.reason}")
|
|
except Exception as e:
|
|
print(f"❌ [Receiver] 连接失败: {e}")
|
|
|
|
async def run_sender():
|
|
"""
|
|
作为发送者连接，同时负责发送音频和接收自己会话的广播结果。
|
|
"""
|
|
await asyncio.sleep(1) # 等待receiver有机会先连接
|
|
print(f"▶️ [Sender] 尝试连接到: {SENDER_URI}")
|
|
try:
|
|
async with websockets.connect(SENDER_URI) as websocket:
|
|
print("✅ [Sender] 连接成功。")
|
|
|
|
# --- 并行任务：接收消息 ---
|
|
async def receive_task():
|
|
print("▶️ [Sender-Receiver] 开始监听广播消息...")
|
|
try:
|
|
while True:
|
|
message = await websocket.recv()
|
|
print(f"🎧 [Sender-Receiver] 收到结果: {message}")
|
|
except websockets.exceptions.ConnectionClosed:
|
|
print("✅ [Sender-Receiver] 连接已关闭，停止监听。")
|
|
|
|
receiver_sub_task = asyncio.create_task(receive_task())
|
|
|
|
# --- 主任务：发送音频 ---
|
|
try:
|
|
print("▶️ [Sender] 准备发送音频...")
|
|
audio_data, sample_rate = sf.read(AUDIO_FILE_PATH, dtype='int16')
|
|
if sample_rate != 16000:
|
|
print(f"❌ [Sender] 错误：音频文件采样率必须是 16kHz。")
|
|
receiver_sub_task.cancel()
|
|
return
|
|
|
|
total_samples = len(audio_data)
|
|
chunk_samples = CHUNK_SIZE // 2
|
|
samples_sent = 0
|
|
print(f"音频加载成功，总长度: {total_samples} samples。开始分块发送...")
|
|
|
|
for i in range(0, total_samples, chunk_samples):
|
|
chunk = audio_data[i:i + chunk_samples]
|
|
if len(chunk) == 0:
|
|
break
|
|
await websocket.send(chunk.tobytes())
|
|
samples_sent += len(chunk)
|
|
print(f"🎧 [Sender] 正在发送: {samples_sent}/{total_samples} samples", end="\r")
|
|
await asyncio.sleep(CHUNK_DURATION_MS / 1000)
|
|
|
|
print()
|
|
print("🏁 [Sender] 音频流发送完毕，发送 'close' 信号。")
|
|
await websocket.send("close")
|
|
|
|
except FileNotFoundError:
|
|
print(f"❌ [Sender] 错误：找不到音频文件 {AUDIO_FILE_PATH}")
|
|
except Exception as e:
|
|
print(f"❌ [Sender] 发送过程中发生错误: {e}")
|
|
|
|
# 等待接收任务自然结束（当连接关闭时）
|
|
await receiver_sub_task
|
|
|
|
except Exception as e:
|
|
print(f"❌ [Sender] 连接失败: {e}")
|
|
|
|
async def main():
|
|
"""同时运行 sender 和 receiver 任务。"""
|
|
print("--- 开始 WebSocket ASR 服务端到端测试 ---")
|
|
print(f"会话 ID: {SESSION_ID}")
|
|
|
|
# 创建 receiver 和 sender 任务
|
|
sender_task = asyncio.create_task(run_sender())
|
|
await asyncio.sleep(7)
|
|
receiver_task = asyncio.create_task(run_receiver())
|
|
|
|
# 等待两个任务完成
|
|
await asyncio.gather(receiver_task, sender_task)
|
|
|
|
print("--- 测试结束 ---")
|
|
|
|
if __name__ == "__main__":
|
|
# 在运行此脚本前，请确保 FastAPI 服务器正在运行。
|
|
# python main.py
|
|
asyncio.run(main())
|