rtcllm/rtcllm/vad.py
2024-09-04 14:37:10 +08:00

147 lines
4.5 KiB
Python

import base64
from traceback import print_exc
import asyncio
import collections
import contextlib
from appPublic.folderUtils import temp_file
from aiortc import MediaStreamTrack
from aiortc.contrib.media import MediaBlackhole, MediaPlayer, MediaRecorder, MediaRelay
import webrtcvad
import wave
import numpy as np
from av import AudioLayout, AudioResampler, AudioFrame, AudioFormat
class AudioTrackVad(MediaStreamTrack):
def __init__(self, track, stage=3, onvoiceend=None):
super().__init__()
self.track = track
print(dir(track), 'AudioTrackVad.__init__()')
self.onvoiceend = onvoiceend
self.vad = webrtcvad.Vad(stage)
# self.sample_rate = self.track.getSettings().sampleRate
# frameSize = self.track.getSettings().frameSize
# self.frame_duration_ms = (1000 * frameSize) / self.sample_rate
self.frame_duration_ms = 0.00002
self.num_padding_frames = 20
self.ring_buffer = collections.deque(maxlen=self.num_padding_frames)
self.triggered = False
self.voiced_frames = []
self.loop = asyncio.get_event_loop()
self.task = None
self.debug = True
self.running = False
def start_vad(self):
self.running = True
self.task = self.loop.call_later(self.frame_duration_ms, self._recv)
def _recv(self):
asyncio.create_task(self.recv())
def stop(self):
self.running = False
def frame2bytes(self, frame):
# 假设你有一个 AudioFrame 对象 audio_frame
audio_array = frame.to_ndarray()
# 将 numpy 数组转换为字节数组
dtype = audio_array.dtype
audio_bytes = audio_array.tobytes()
return audio_bytes
async def recv(self):
oldf = await self.track.recv()
self.sample_rate = oldf.sample_rate
frames = self.resample(oldf)
for f in frames:
if self.debug:
self.debug = False
print(f'{type(f)}, {f.samples=}, {f.format.bytes=}, {f.sample_rate=}, {f.format=}, {f.is_corrupt=}, {f.layout=}, {f.planes=}, {f.side_data=}')
try:
await self.vad_check(f)
except Exception as e:
print(f'{e=}')
print_exc()
return
if self.task:
self.task.cancel()
if self.running:
self.task = self.loop.call_later(self.frame_duration_ms, self._recv)
return f
def resample(self, frame, sample_rate=None):
if sample_rate is None:
sample_rate = frame.rate
r = AudioResampler(format='s16', layout='mono', rate=sample_rate)
frame = r.resample(frame)
return frame
async def vad_check(self, frame):
is_speech = self.vad.is_speech(self.frame2bytes(frame), self.sample_rate)
if not self.triggered:
self.ring_buffer.append((frame, is_speech))
num_voiced = len([f for f, speech in self.ring_buffer if speech])
# If we're NOTTRIGGERED and more than 90% of the frames in
# the ring buffer are voiced frames, then enter the
# TRIGGERED state.
if num_voiced > 0.9 * self.ring_buffer.maxlen:
self.triggered = True
# We want to yield all the audio we see from now until
# we are NOTTRIGGERED, but we have to start with the
# audio that's already in the ring buffer.
for f, s in self.ring_buffer:
self.voiced_frames.append(f)
self.ring_buffer.clear()
print('start voice .....', len(self.voiced_frames))
else:
# We're in the TRIGGERED state, so collect the audio data
# and add it to the ring buffer.
self.voiced_frames.append(frame)
self.ring_buffer.append((frame, is_speech))
num_unvoiced = len([f for f, speech in self.ring_buffer if not speech])
# If more than 90% of the frames in the ring buffer are
# unvoiced, then enter NOTTRIGGERED and yield whatever
# audio we've collected.
if num_unvoiced > 0.9 * self.ring_buffer.maxlen:
self.triggered = False
ret = await self.write_wave()
# ret = await self.gen_base64()
if self.onvoiceend:
await self.onvoiceend(ret)
self.ring_buffer.clear()
self.voiced_frames = []
def to_mono16000_data(self):
lst = []
for f in self.voiced_frames:
fs = self.resample(f, sample_rate=16000)
lst += fs
audio_data = b''.join([self.frame2bytes(f) for f in lst])
return audio_data
async def gen_base64(self):
audio_data = self.to_mono16000_data()
b64 = base64.b64encode(audio_data).decode('utf-8')
return b64
async def write_wave(self):
"""Writes a .wav file.
Takes path, PCM audio data, and sample rate.
"""
audio_data = self.to_mono16000_data()
path = temp_file(suffix='.wav')
print(f'temp_file={path}')
with contextlib.closing(wave.open(path, 'wb')) as wf:
wf.setnchannels(1)
wf.setsampwidth(2)
wf.setframerate(16000)
wf.writeframes(audio_data)
print('************wrote*******')
if self.onvoiceend:
await self.onvoiceend(path)
print('************over*******')