97 lines
3.0 KiB
Python
97 lines
3.0 KiB
Python
from fastapi import FastAPI, WebSocket, WebSocketDisconnect
|
||
import uvicorn
|
||
from starlette.middleware.cors import CORSMiddleware
|
||
import wave
|
||
import io
|
||
import asyncio
|
||
|
||
app = FastAPI()
|
||
|
||
# 跨域配置
|
||
app.add_middleware(
|
||
CORSMiddleware,
|
||
allow_origins=["*"],
|
||
allow_credentials=True,
|
||
allow_methods=["*"],
|
||
allow_headers=["*"],
|
||
)
|
||
|
||
# 存储活跃连接
|
||
active_connections = []
|
||
|
||
# 音频参数（必须和前端一致！）
|
||
AUDIO_PARAMS = {
|
||
"sample_rate": 16000, # 前端设置的 16000
|
||
"channels": 1, # 单声道
|
||
"sample_width": 2, # 16位深度（2字节/样本）
|
||
"frame_size": 1024 # 前端设置的 1024
|
||
}
|
||
|
||
|
||
def pcm_to_wav(pcm_data: bytes, params: dict) -> bytes:
|
||
"""将裸 PCM 数据封装成 WAV 格式（带文件头）"""
|
||
output = io.BytesIO()
|
||
with wave.open(output, 'wb') as wf:
|
||
wf.setnchannels(params["channels"])
|
||
wf.setsampwidth(params["sample_width"])
|
||
wf.setframerate(params["sample_rate"])
|
||
wf.writeframes(pcm_data)
|
||
# 重置指针并返回字节数据
|
||
output.seek(0)
|
||
return output.read()
|
||
|
||
|
||
async def broadcast_audio(audio_data: bytes):
|
||
"""广播音频数据给所有连接的客户端"""
|
||
for connection in active_connections:
|
||
try:
|
||
await connection.send_bytes(audio_data)
|
||
except Exception as e:
|
||
print(f"广播失败：{e}")
|
||
|
||
|
||
@app.websocket("/ws/audio")
|
||
async def websocket_audio(websocket: WebSocket):
|
||
await websocket.accept()
|
||
active_connections.append(websocket)
|
||
|
||
# 可选：后端本地播放测试（需安装 pyaudio）
|
||
import pyaudio
|
||
p = pyaudio.PyAudio()
|
||
stream = p.open(format=p.get_format_from_width(AUDIO_PARAMS["sample_width"]),
|
||
channels=AUDIO_PARAMS["channels"],
|
||
rate=AUDIO_PARAMS["sample_rate"],
|
||
output=True)
|
||
|
||
try:
|
||
while True:
|
||
# 1. 接收前端的裸 PCM 音频帧
|
||
pcm_data = await websocket.receive_bytes()
|
||
# print(f"收到 PCM 帧：{len(pcm_data)} bytes")
|
||
|
||
# 2. 可选1：后端本地实时播放（测试用）
|
||
stream.write(pcm_data)
|
||
|
||
# 3. 可选2：封装成 WAV 并回传给前端（前端可直接播放）
|
||
# wav_data = pcm_to_wav(pcm_data, AUDIO_PARAMS)
|
||
# await websocket.send_bytes(wav_data) # 仅返回给当前客户端
|
||
# await broadcast_audio(wav_data) # 广播给所有客户端
|
||
|
||
except WebSocketDisconnect:
|
||
active_connections.remove(websocket)
|
||
print(f"客户端断开，当前连接数：{len(active_connections)}")
|
||
except Exception as e:
|
||
if websocket in active_connections:
|
||
active_connections.remove(websocket)
|
||
print(f"处理音频失败：{e}")
|
||
finally:
|
||
# 关闭本地播放流
|
||
stream.stop_stream()
|
||
stream.close()
|
||
p.terminate()
|
||
|
||
|
||
if __name__ == "__main__":
|
||
# 启动服务（确保安装了 websockets：pip install uvicorn[standard]）
|
||
uvicorn.run(app, host="0.0.0.0", port=8000)
|