websocket online vad endpoint

游雁

2023-05-17 ed2d1fed0541b3e097683f07571cda132f9c06c5

websocket online vad endpoint

 funasr/runtime/python/websocket/ws_server_online.py

@@ -32,15 +32,29 @@
    ncpu=args.ncpu,
    model_revision='v1.0.4')

# vad
inference_pipeline_vad = pipeline(
    task=Tasks.voice_activity_detection,
    model=args.vad_model,
    model_revision=None,
    output_dir=None,
    batch_size=1,
    mode='online',
    ngpu=args.ngpu,
    ncpu=1,
)

print("model loaded")



async def ws_serve(websocket, path):
    frames = []
    frames_asr_online = []
    global websocket_users
    websocket_users.add(websocket)
    websocket.param_dict_asr_online = {"cache": dict()}
    websocket.param_dict_vad = {'in_cache': dict()}
    websocket.wav_name = "microphone"
    print("new user connected",flush=True)
    try:
@@ -53,6 +67,7 @@
                if "is_speaking" in messagejson:
                    websocket.is_speaking = messagejson["is_speaking"]
                    websocket.param_dict_asr_online["is_final"] = not websocket.is_speaking
                    websocket.param_dict_vad["is_final"] = not websocket.is_speaking
                    # need to fire engine manually if no data received any more
                    if not websocket.is_speaking:
                        await async_asr_online(websocket,b"")
@@ -66,11 +81,15 @@
            if len(frames_asr_online) > 0 or not isinstance(message, str):
                if not isinstance(message,str):
                    frames_asr_online.append(message)
                    # frames.append(message)
                    # duration_ms = len(message) // 32
                    # websocket.vad_pre_idx += duration_ms
                    speech_start_i, speech_end_i = await async_vad(websocket, message)
                    websocket.is_speaking = not speech_end_i
					
                if len(frames_asr_online) % websocket.chunk_interval == 0 or not websocket.is_speaking:
                    websocket.param_dict_asr_online["is_final"] = not websocket.is_speaking
                    audio_in = b"".join(frames_asr_online)
                    # if not websocket.is_speaking:
                        #padding 0.5s at end gurantee that asr engine can fire out last word
                        # audio_in=audio_in+b''.join(np.zeros(int(16000*0.5),dtype=np.int16))
                    await async_asr_online(websocket,audio_in)
                    frames_asr_online = []
    
@@ -97,6 +116,20 @@
                await websocket.send(message)


async def async_vad(websocket, audio_in):
    segments_result = inference_pipeline_vad(audio_in=audio_in, param_dict=websocket.param_dict_vad)
	
    speech_start = False
    speech_end = False
	
    if len(segments_result) == 0 or len(segments_result["text"]) > 1:
        return speech_start, speech_end
    if segments_result["text"][0][0] != -1:
        speech_start = segments_result["text"][0][0]
    if segments_result["text"][0][1] != -1:
        speech_end = True
    return speech_start, speech_end

if len(args.certfile)>0:
  ssl_context = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER)