77 lines
2.6 KiB
Python
77 lines
2.6 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
from fractions import Fraction
|
|
from collections import deque
|
|
|
|
import numpy as np
|
|
from aiortc import MediaStreamTrack
|
|
from av import AudioFrame
|
|
|
|
|
|
class AudioBus:
|
|
def __init__(self, sample_rate: int = 48000) -> None:
|
|
self.sample_rate = sample_rate
|
|
self._chunks: deque[np.ndarray] = deque()
|
|
self._lock = asyncio.Lock()
|
|
self._level = 0.0
|
|
|
|
async def enqueue(self, pcm_int16_mono: np.ndarray) -> None:
|
|
async with self._lock:
|
|
if pcm_int16_mono.ndim != 1:
|
|
pcm_int16_mono = pcm_int16_mono.reshape(-1)
|
|
chunk = pcm_int16_mono.astype(np.int16, copy=False)
|
|
self._chunks.append(chunk)
|
|
if chunk.size > 0:
|
|
rms = float(np.sqrt(np.mean((chunk.astype(np.float32) / 32768.0) ** 2)))
|
|
self._level = 0.85 * self._level + 0.15 * min(1.0, rms * 10.0)
|
|
|
|
async def read(self, samples: int) -> np.ndarray:
|
|
out = np.zeros(samples, dtype=np.int16)
|
|
async with self._lock:
|
|
i = 0
|
|
while i < samples and self._chunks:
|
|
head = self._chunks[0]
|
|
n = min(samples - i, head.shape[0])
|
|
out[i : i + n] = head[:n]
|
|
i += n
|
|
if n == head.shape[0]:
|
|
self._chunks.popleft()
|
|
else:
|
|
self._chunks[0] = head[n:]
|
|
self._level = 0.92 * self._level
|
|
return out
|
|
|
|
async def level(self) -> float:
|
|
async with self._lock:
|
|
return float(self._level)
|
|
|
|
async def clear(self) -> None:
|
|
async with self._lock:
|
|
self._chunks.clear()
|
|
self._level = 0.0
|
|
|
|
|
|
class AvatarAudioTrack(MediaStreamTrack):
|
|
kind = "audio"
|
|
|
|
def __init__(self, audio_bus: AudioBus, sample_rate: int = 48000, frame_ms: int = 20) -> None:
|
|
super().__init__()
|
|
self.audio_bus = audio_bus
|
|
self.sample_rate = sample_rate
|
|
self.samples_per_frame = int(sample_rate * frame_ms / 1000)
|
|
self._pts = 0
|
|
self._time_base = Fraction(1, sample_rate)
|
|
|
|
async def recv(self) -> AudioFrame:
|
|
await asyncio.sleep(self.samples_per_frame / self.sample_rate)
|
|
mono = await self.audio_bus.read(self.samples_per_frame)
|
|
samples = mono.reshape(1, -1)
|
|
frame = AudioFrame(format="s16", layout="mono", samples=self.samples_per_frame)
|
|
frame.planes[0].update(samples.tobytes())
|
|
frame.sample_rate = self.sample_rate
|
|
frame.pts = self._pts
|
|
frame.time_base = self._time_base
|
|
self._pts += self.samples_per_frame
|
|
return frame
|