412 lines
16 KiB
Python
412 lines
16 KiB
Python
#!/usr/bin/env python3
|
||
# -*- coding: utf-8 -*-
|
||
|
||
"""
|
||
豆包音频处理模块 - 最终简化版本
|
||
实现音频文件上传和TTS音频下载的完整流程
|
||
"""
|
||
|
||
import asyncio
|
||
import gzip
|
||
import json
|
||
import uuid
|
||
import wave
|
||
import time
|
||
import os
|
||
from typing import Dict, Any, Optional
|
||
import websockets
|
||
|
||
|
||
# 协议常量
|
||
PROTOCOL_VERSION = 0b0001
|
||
CLIENT_FULL_REQUEST = 0b0001
|
||
CLIENT_AUDIO_ONLY_REQUEST = 0b0010
|
||
SERVER_FULL_RESPONSE = 0b1001
|
||
SERVER_ACK = 0b1011
|
||
SERVER_ERROR_RESPONSE = 0b1111
|
||
|
||
NO_SEQUENCE = 0b0000
|
||
MSG_WITH_EVENT = 0b0100
|
||
|
||
NO_SERIALIZATION = 0b0000
|
||
JSON = 0b0001
|
||
GZIP = 0b0001
|
||
|
||
|
||
def generate_header(
|
||
version=PROTOCOL_VERSION,
|
||
message_type=CLIENT_FULL_REQUEST,
|
||
message_type_specific_flags=MSG_WITH_EVENT,
|
||
serial_method=JSON,
|
||
compression_type=GZIP,
|
||
reserved_data=0x00,
|
||
extension_header=bytes()
|
||
):
|
||
"""生成协议头"""
|
||
header = bytearray()
|
||
header_size = int(len(extension_header) / 4) + 1
|
||
header.append((version << 4) | header_size)
|
||
header.append((message_type << 4) | message_type_specific_flags)
|
||
header.append((serial_method << 4) | compression_type)
|
||
header.append(reserved_data)
|
||
header.extend(extension_header)
|
||
return header
|
||
|
||
|
||
class DoubaoClient:
|
||
"""豆包客户端"""
|
||
|
||
def __init__(self):
|
||
self.base_url = "wss://openspeech.bytedance.com/api/v3/realtime/dialogue"
|
||
self.app_id = "8718217928"
|
||
self.access_key = "ynJMX-5ix1FsJvswC9KTNlGUdubcchqc"
|
||
self.app_key = "PlgvMymc7f3tQnJ6"
|
||
self.resource_id = "volc.speech.dialog"
|
||
self.session_id = str(uuid.uuid4())
|
||
self.ws = None
|
||
self.log_id = ""
|
||
|
||
def get_headers(self) -> Dict[str, str]:
|
||
"""获取请求头"""
|
||
return {
|
||
"X-Api-App-ID": self.app_id,
|
||
"X-Api-Access-Key": self.access_key,
|
||
"X-Api-Resource-Id": self.resource_id,
|
||
"X-Api-App-Key": self.app_key,
|
||
"X-Api-Connect-Id": str(uuid.uuid4()),
|
||
}
|
||
|
||
async def connect(self) -> None:
|
||
"""建立WebSocket连接"""
|
||
print(f"连接豆包服务器: {self.base_url}")
|
||
|
||
try:
|
||
self.ws = await websockets.connect(
|
||
self.base_url,
|
||
additional_headers=self.get_headers(),
|
||
ping_interval=None
|
||
)
|
||
|
||
# 获取log_id
|
||
if hasattr(self.ws, 'response_headers'):
|
||
self.log_id = self.ws.response_headers.get("X-Tt-Logid")
|
||
elif hasattr(self.ws, 'headers'):
|
||
self.log_id = self.ws.headers.get("X-Tt-Logid")
|
||
|
||
print(f"连接成功, log_id: {self.log_id}")
|
||
|
||
# 发送StartConnection请求
|
||
await self._send_start_connection()
|
||
|
||
# 发送StartSession请求
|
||
await self._send_start_session()
|
||
|
||
except Exception as e:
|
||
print(f"连接失败: {e}")
|
||
raise
|
||
|
||
async def _send_start_connection(self) -> None:
|
||
"""发送StartConnection请求"""
|
||
print("发送StartConnection请求...")
|
||
request = bytearray(generate_header())
|
||
request.extend(int(1).to_bytes(4, 'big'))
|
||
|
||
payload_bytes = b"{}"
|
||
payload_bytes = gzip.compress(payload_bytes)
|
||
request.extend(len(payload_bytes).to_bytes(4, 'big'))
|
||
request.extend(payload_bytes)
|
||
|
||
await self.ws.send(request)
|
||
response = await self.ws.recv()
|
||
print(f"StartConnection响应长度: {len(response)}")
|
||
|
||
async def _send_start_session(self) -> None:
|
||
"""发送StartSession请求"""
|
||
print("发送StartSession请求...")
|
||
session_config = {
|
||
"asr": {
|
||
"extra": {
|
||
"end_smooth_window_ms": 1500,
|
||
},
|
||
},
|
||
"tts": {
|
||
"speaker": "zh_female_vv_jupiter_bigtts",
|
||
"audio_config": {
|
||
"channel": 1,
|
||
"format": "pcm",
|
||
"sample_rate": 24000
|
||
},
|
||
},
|
||
"dialog": {
|
||
"bot_name": "豆包",
|
||
"system_role": "你使用活泼灵动的女声，性格开朗，热爱生活。",
|
||
"speaking_style": "你的说话风格简洁明了，语速适中，语调自然。",
|
||
"location": {"city": "北京"},
|
||
"extra": {
|
||
"strict_audit": False,
|
||
"audit_response": "支持客户自定义安全审核回复话术。",
|
||
"recv_timeout": 30,
|
||
"input_mod": "audio",
|
||
},
|
||
},
|
||
}
|
||
|
||
request = bytearray(generate_header())
|
||
request.extend(int(100).to_bytes(4, 'big'))
|
||
request.extend(len(self.session_id).to_bytes(4, 'big'))
|
||
request.extend(self.session_id.encode())
|
||
|
||
payload_bytes = json.dumps(session_config).encode()
|
||
payload_bytes = gzip.compress(payload_bytes)
|
||
request.extend(len(payload_bytes).to_bytes(4, 'big'))
|
||
request.extend(payload_bytes)
|
||
|
||
await self.ws.send(request)
|
||
response = await self.ws.recv()
|
||
print(f"StartSession响应长度: {len(response)}")
|
||
|
||
# 等待一会确保会话完全建立
|
||
await asyncio.sleep(1.0)
|
||
|
||
async def task_request(self, audio: bytes) -> None:
|
||
"""发送音频数据"""
|
||
task_request = bytearray(
|
||
generate_header(message_type=CLIENT_AUDIO_ONLY_REQUEST,
|
||
serial_method=NO_SERIALIZATION))
|
||
task_request.extend(int(200).to_bytes(4, 'big'))
|
||
task_request.extend((len(self.session_id)).to_bytes(4, 'big'))
|
||
task_request.extend(str.encode(self.session_id))
|
||
payload_bytes = gzip.compress(audio)
|
||
task_request.extend((len(payload_bytes)).to_bytes(4, 'big'))
|
||
task_request.extend(payload_bytes)
|
||
await self.ws.send(task_request)
|
||
|
||
def parse_response(self, response):
|
||
"""解析响应"""
|
||
if len(response) < 4:
|
||
return None
|
||
|
||
protocol_version = response[0] >> 4
|
||
header_size = response[0] & 0x0f
|
||
message_type = response[1] >> 4
|
||
flags = response[1] & 0x0f
|
||
|
||
payload_start = header_size * 4
|
||
payload = response[payload_start:]
|
||
|
||
result = {
|
||
'protocol_version': protocol_version,
|
||
'header_size': header_size,
|
||
'message_type': message_type,
|
||
'flags': flags,
|
||
'payload': payload,
|
||
'payload_size': len(payload)
|
||
}
|
||
|
||
# 解析payload
|
||
if len(payload) >= 4:
|
||
result['event'] = int.from_bytes(payload[:4], 'big')
|
||
|
||
if len(payload) >= 8:
|
||
session_id_len = int.from_bytes(payload[4:8], 'big')
|
||
if len(payload) >= 8 + session_id_len:
|
||
result['session_id'] = payload[8:8+session_id_len].decode()
|
||
|
||
if len(payload) >= 12 + session_id_len:
|
||
data_size = int.from_bytes(payload[8+session_id_len:12+session_id_len], 'big')
|
||
result['data_size'] = data_size
|
||
result['data'] = payload[12+session_id_len:12+session_id_len+data_size]
|
||
|
||
# 尝试解析JSON数据
|
||
try:
|
||
result['json_data'] = json.loads(result['data'].decode('utf-8'))
|
||
except:
|
||
pass
|
||
|
||
return result
|
||
|
||
async def process_audio_file(self, input_file: str, output_file: str) -> bool:
|
||
"""处理音频文件：上传并获得TTS响应"""
|
||
print(f"开始处理音频文件: {input_file}")
|
||
|
||
try:
|
||
# 读取输入音频文件
|
||
with wave.open(input_file, 'rb') as wf:
|
||
audio_data = wf.readframes(wf.getnframes())
|
||
print(f"读取音频数据: {len(audio_data)} 字节")
|
||
print(f"音频参数: {wf.getframerate()}Hz, {wf.getnchannels()}通道, {wf.getsampwidth()*8}-bit")
|
||
|
||
# 发送音频数据
|
||
print("发送音频数据...")
|
||
await self.task_request(audio_data)
|
||
print("音频数据发送成功")
|
||
|
||
# 接收响应序列
|
||
print("开始接收响应...")
|
||
audio_chunks = []
|
||
response_count = 0
|
||
max_responses = 20
|
||
|
||
while response_count < max_responses:
|
||
try:
|
||
response = await asyncio.wait_for(self.ws.recv(), timeout=30.0)
|
||
response_count += 1
|
||
|
||
parsed = self.parse_response(response)
|
||
if not parsed:
|
||
continue
|
||
|
||
print(f"响应 {response_count}: message_type={parsed['message_type']}, event={parsed.get('event', 'N/A')}, size={parsed['payload_size']}")
|
||
|
||
# 处理不同类型的响应
|
||
if parsed['message_type'] == 11: # SERVER_ACK - 可能包含音频
|
||
if 'data' in parsed and parsed['data_size'] > 0:
|
||
audio_chunks.append(parsed['data'])
|
||
print(f"收集到音频块: {parsed['data_size']} 字节")
|
||
|
||
elif parsed['message_type'] == 9: # SERVER_FULL_RESPONSE
|
||
event = parsed.get('event', 0)
|
||
|
||
if event == 350: # TTS开始
|
||
print("TTS音频生成开始")
|
||
elif event == 359: # TTS结束
|
||
print("TTS音频生成结束")
|
||
break
|
||
elif event == 451: # ASR结果
|
||
if 'json_data' in parsed and 'results' in parsed['json_data']:
|
||
text = parsed['json_data']['results'][0].get('text', '')
|
||
print(f"语音识别结果: {text}")
|
||
elif event == 550: # TTS音频数据
|
||
if 'data' in parsed and parsed['data_size'] > 0:
|
||
# 检查是否是JSON（音频元数据）还是实际音频数据
|
||
try:
|
||
json.loads(parsed['data'].decode('utf-8'))
|
||
print("收到TTS音频元数据")
|
||
except:
|
||
# 不是JSON，可能是音频数据
|
||
audio_chunks.append(parsed['data'])
|
||
print(f"收集到TTS音频块: {parsed['data_size']} 字节")
|
||
|
||
except asyncio.TimeoutError:
|
||
print(f"等待响应 {response_count + 1} 超时")
|
||
break
|
||
except websockets.exceptions.ConnectionClosed:
|
||
print("连接已关闭")
|
||
break
|
||
|
||
print(f"共收到 {response_count} 个响应，收集到 {len(audio_chunks)} 个音频块")
|
||
|
||
# 合并音频数据
|
||
if audio_chunks:
|
||
combined_audio = b''.join(audio_chunks)
|
||
print(f"合并后的音频数据: {len(combined_audio)} 字节")
|
||
|
||
# 检查是否是GZIP压缩数据
|
||
try:
|
||
decompressed = gzip.decompress(combined_audio)
|
||
print(f"解压缩后音频数据: {len(decompressed)} 字节")
|
||
audio_to_write = decompressed
|
||
except:
|
||
print("音频数据不是GZIP压缩格式，直接使用原始数据")
|
||
audio_to_write = combined_audio
|
||
|
||
# 创建输出WAV文件
|
||
try:
|
||
# 豆包返回的音频是32位浮点格式，需要转换为16位整数
|
||
import struct
|
||
|
||
# 检查音频数据长度是否是4的倍数（32位浮点）
|
||
if len(audio_to_write) % 4 != 0:
|
||
print(f"警告：音频数据长度 {len(audio_to_write)} 不是4的倍数，截断到最近的倍数")
|
||
audio_to_write = audio_to_write[:len(audio_to_write) // 4 * 4]
|
||
|
||
# 将32位浮点转换为16位整数
|
||
float_count = len(audio_to_write) // 4
|
||
int16_data = bytearray(float_count * 2)
|
||
|
||
for i in range(float_count):
|
||
# 读取32位浮点数（小端序）
|
||
float_value = struct.unpack('<f', audio_to_write[i*4:i*4+4])[0]
|
||
|
||
# 将浮点数限制在[-1.0, 1.0]范围内
|
||
float_value = max(-1.0, min(1.0, float_value))
|
||
|
||
# 转换为16位整数
|
||
int16_value = int(float_value * 32767)
|
||
|
||
# 写入16位整数（小端序）
|
||
int16_data[i*2:i*2+2] = struct.pack('<h', int16_value)
|
||
|
||
# 创建WAV文件
|
||
with wave.open(output_file, 'wb') as wav_file:
|
||
wav_file.setnchannels(1)
|
||
wav_file.setsampwidth(2)
|
||
wav_file.setframerate(24000)
|
||
wav_file.writeframes(int16_data)
|
||
|
||
print(f"成功创建输出文件: {output_file}")
|
||
print(f"音频转换: {float_count} 个32位浮点样本 -> {len(int16_data)//2} 个16位整数样本")
|
||
|
||
# 显示文件信息
|
||
if os.path.exists(output_file):
|
||
file_size = os.path.getsize(output_file)
|
||
duration = file_size / (24000 * 1 * 2)
|
||
print(f"输出文件大小: {file_size} 字节，时长: {duration:.2f} 秒")
|
||
|
||
return True
|
||
|
||
except Exception as e:
|
||
print(f"创建WAV文件失败: {e}")
|
||
# 保存原始数据
|
||
with open(output_file + '.raw', 'wb') as f:
|
||
f.write(audio_to_write)
|
||
print(f"原始音频数据已保存到: {output_file}.raw")
|
||
return False
|
||
else:
|
||
print("未收到音频数据")
|
||
return False
|
||
|
||
except Exception as e:
|
||
print(f"处理音频文件失败: {e}")
|
||
import traceback
|
||
traceback.print_exc()
|
||
return False
|
||
|
||
async def close(self) -> None:
|
||
"""关闭连接"""
|
||
if self.ws:
|
||
try:
|
||
await self.ws.close()
|
||
except:
|
||
pass
|
||
print("连接已关闭")
|
||
|
||
|
||
async def main():
|
||
"""主函数"""
|
||
client = DoubaoClient()
|
||
|
||
try:
|
||
await client.connect()
|
||
|
||
# 处理录音文件
|
||
input_file = "recording_20250920_135137.wav"
|
||
output_file = "tts_output.wav"
|
||
|
||
success = await client.process_audio_file(input_file, output_file)
|
||
|
||
if success:
|
||
print("音频处理成功！")
|
||
else:
|
||
print("音频处理失败")
|
||
|
||
except Exception as e:
|
||
print(f"程序失败: {e}")
|
||
import traceback
|
||
traceback.print_exc()
|
||
finally:
|
||
await client.close()
|
||
|
||
|
||
if __name__ == "__main__":
|
||
asyncio.run(main()) |