init

webrtc/tracks.py

@@ -0,0 +1,114 @@
+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, VideoFrame
+
+from core.state_machine import Arbitrator, SessionState
+from services.avatar import AvatarService
+
+
+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 AvatarVideoTrack(MediaStreamTrack):
+    kind = "video"
+
+    def __init__(
+        self,
+        avatar_service: AvatarService,
+        arbitrator: Arbitrator,
+        audio_bus: AudioBus,
+        fps: int = 25,
+    ) -> None:
+        super().__init__()
+        self.avatar_service = avatar_service
+        self.arbitrator = arbitrator
+        self.audio_bus = audio_bus
+        self.fps = fps
+        self._pts = 0
+        self._time_base = Fraction(1, 90000)
+
+    async def recv(self) -> VideoFrame:
+        await asyncio.sleep(1 / self.fps)
+        generated = await self.avatar_service.pop_generated_frame()
+        if generated is not None:
+            arr = generated
+        else:
+            level = await self.audio_bus.level()
+            speaking = self.arbitrator.state == SessionState.AVATAR_SPEAKING or level > 0.02
+            arr = await self.avatar_service.render_frame(mouth_open=level, speaking=speaking)
+
+        frame = VideoFrame.from_ndarray(arr, format="bgr24")
+        frame.pts = self._pts
+        frame.time_base = self._time_base
+        self._pts += int(90000 / self.fps)
+        return frame
+
+
+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