131 lines
4.2 KiB
Python
131 lines
4.2 KiB
Python
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.00008
|
|
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()
|
|
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=}')
|
|
self.sample_rate = f.sample_rate
|
|
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):
|
|
fmt = AudioFormat('s16')
|
|
al = AudioLayout(1)
|
|
r = AudioResampler(format=fmt, layout=al, rate=frame.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
|
|
audio_data = b''.join([self.frame2bytes(f) for f in self.voiced_frames])
|
|
await self.write_wave(audio_data)
|
|
self.ring_buffer.clear()
|
|
self.voiced_frames = []
|
|
print('end voice .....', len(self.voiced_frames))
|
|
|
|
async def write_wave(self, audio_data):
|
|
"""Writes a .wav file.
|
|
|
|
Takes path, PCM audio data, and sample rate.
|
|
"""
|
|
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(self.sample_rate)
|
|
wf.writeframes(audio_data)
|
|
|
|
print('************wrote*******')
|
|
if self.onvoiceend:
|
|
await self.onvoiceend(path)
|
|
print('************over*******')
|
|
|