196 lines
6.9 KiB
Python
196 lines
6.9 KiB
Python
import json, os, threading, time, traceback
|
||
from typing import (
|
||
Any, List
|
||
)
|
||
from copy import deepcopy
|
||
from utils.logger import logger
|
||
|
||
from websocket import (
|
||
create_connection,
|
||
WebSocketConnectionClosedException,
|
||
ABNF
|
||
)
|
||
|
||
from utils.model import (
|
||
ASRResponseModel,
|
||
SegmentModel
|
||
)
|
||
import threading
|
||
import queue
|
||
|
||
from pydantic_core import ValidationError
|
||
|
||
_IS_TEST = os.getenv("SUBMIT_CONFIG_FILEPATH") is None
|
||
|
||
|
||
class ASRWebSocketClient:
|
||
def __init__(self, url: str):
|
||
self.endpoint = f"{url}/recognition"
|
||
# self.ctx = deepcopy(ctx)
|
||
self.ws = None
|
||
self.conn_attempts = -1
|
||
self.failed = False
|
||
self.terminate_time = float("inf")
|
||
self.sent_timestamps: List[float] = []
|
||
self.received_timestamps: List[float] = []
|
||
self.results: List[Any] = []
|
||
self.connected = True
|
||
logger.info(f"Target endpoint: {self.endpoint}")
|
||
|
||
def execute(self, path: str) -> List[SegmentModel]:
|
||
# 开启线程，一个发，一个接要记录接收数据的时间。
|
||
send_thread = threading.Thread(target=self.send, args=(path,))
|
||
# 记录开始时间
|
||
start_time = time.time()
|
||
# 启动线程发送数据
|
||
send_thread.start()
|
||
|
||
# 用来返回结果的队列
|
||
result_queue = queue.Queue()
|
||
|
||
# 线程封装 receive
|
||
def receive_thread_fn():
|
||
try:
|
||
res = self.receive(start_time)
|
||
result_queue.put(res)
|
||
except Exception as e:
|
||
# 放异常信息
|
||
result_queue.put(None)
|
||
|
||
# 启动 receive 线程（关注返回值）
|
||
receive_thread = threading.Thread(target=receive_thread_fn)
|
||
receive_thread.start()
|
||
|
||
# 可选：等待 send 线程也结束（可删）
|
||
send_thread.join()
|
||
receive_thread.join()
|
||
|
||
# 等待 receive 完成，并获取返回值
|
||
result = result_queue.get()
|
||
|
||
return result
|
||
|
||
def initialize_connection(self, language: str):
|
||
expiration = time.time() + float(os.getenv("end_time", "2"))
|
||
self.connected = False
|
||
init = False
|
||
while True:
|
||
try:
|
||
if not init:
|
||
logger.debug(f"建立ws链接：发送建立ws链接请求。")
|
||
self.ws = create_connection(self.endpoint)
|
||
body = json.dumps(self._get_init_payload(language))
|
||
logger.debug(f"建立ws链接：发送初始化数据。{body}")
|
||
self.ws.send(body)
|
||
init = True
|
||
msg = self.ws.recv()
|
||
logger.debug(f"收到响应数据: {msg}")
|
||
if len(msg) == 0:
|
||
time.sleep(0.5) # 睡眠一下等待数据写回来
|
||
continue
|
||
if isinstance(msg, str):
|
||
try:
|
||
msg = json.loads(msg)
|
||
except Exception:
|
||
raise Exception("建立ws链接：响应数据非json格式！")
|
||
if isinstance(msg, dict):
|
||
connected = msg.get("success")
|
||
if connected:
|
||
logger.debug("建立ws链接：链接建立成功！")
|
||
self.conn_attempts = 0
|
||
self.connected = True
|
||
return
|
||
else:
|
||
logger.info("建立ws链接：链接建立失败！")
|
||
init = False
|
||
self.conn_attempts = self.conn_attempts + 1
|
||
if self.conn_attempts > 5:
|
||
raise ConnectionRefusedError("重试5次后，仍然无法建立ws链接。")
|
||
if time.time() > expiration:
|
||
raise RuntimeError("建立ws链接：链接建立超时！")
|
||
|
||
except WebSocketConnectionClosedException or TimeoutError:
|
||
raise Exception("建立ws链接：初始化阶段连接中断，退出。")
|
||
except Exception as e:
|
||
logger.info("建立ws链接：链接建立失败！")
|
||
init = False
|
||
self.conn_attempts = self.conn_attempts + 1
|
||
if self.conn_attempts > 5:
|
||
raise ConnectionRefusedError("重试5次后，仍然无法建立ws链接。")
|
||
if time.time() > expiration:
|
||
raise RuntimeError("建立ws链接：链接建立超时！")
|
||
|
||
def shutdown(self):
|
||
try:
|
||
if self.ws:
|
||
self.ws.close()
|
||
except Exception:
|
||
pass
|
||
self.connected = False
|
||
|
||
def _get_init_payload(self, language="zh") -> dict:
|
||
# language = "zh"
|
||
return {
|
||
"language": language
|
||
}
|
||
|
||
def _get_finish_payload(self) -> dict:
|
||
return {
|
||
"end": "true"
|
||
}
|
||
|
||
def send(self, path):
|
||
skip_wav_header = path.endswith("wav")
|
||
with open(path, "rb") as f:
|
||
if skip_wav_header:
|
||
# WAV 文件头部为 44 字节
|
||
f.read(44)
|
||
|
||
while chunk := f.read(3200):
|
||
logger.debug(f"发送 {len(chunk)} 字节数据.")
|
||
self.ws.send(chunk, opcode=ABNF.OPCODE_BINARY)
|
||
time.sleep(0.1)
|
||
self.ws.send(json.dumps(self._get_finish_payload()))
|
||
|
||
def receive(self, start_time) -> [SegmentModel]:
|
||
results = []
|
||
while True:
|
||
msg = self.ws.recv()
|
||
# 记录读取到数据的时间（毫秒值）
|
||
now = 1000 * (time.time() - start_time)
|
||
logger.debug(f"{now} 收到响应数据: {msg}")
|
||
res = json.loads(msg)
|
||
if res.get("asr_results"):
|
||
item = ASRResponseModel.model_validate_json(msg).asr_results
|
||
item.receive_time = now
|
||
results.append(item)
|
||
# logger.info(item.summary())
|
||
else:
|
||
logger.debug(f"响应结束")
|
||
# 按照para_seq排序，并检查一下序号是否连续
|
||
results.sort(key=lambda x: x.para_seq)
|
||
|
||
missing_seqs = []
|
||
for i in range(len(results) - 1):
|
||
expected_next = results[i].para_seq + 1
|
||
actual_next = results[i + 1].para_seq
|
||
if actual_next != expected_next:
|
||
missing_seqs.extend(range(expected_next, actual_next))
|
||
|
||
if missing_seqs:
|
||
logger.warning(f"检测到丢失的段落序号：{missing_seqs}")
|
||
else:
|
||
logger.debug("响应数据正常")
|
||
|
||
return results
|
||
|
||
|
||
if __name__ == '__main__':
|
||
ws_client = ASRWebSocketClient("ws://localhost:18000")
|
||
ws_client.initialize_connection("zh")
|
||
|
||
res = ws_client.execute("/Users/yu/Documents/code-work/asr-live-iluvatar/zh_250312/zh/99.wav")
|
||
for i in res:
|
||
print(i.summary())
|
||
print()
|