AudioStream play example with noise cancellation (#467)

Author: typester

examples/play_audio_stream.py

@@ -0,0 +1,158 @@
+import asyncio
+import os
+import numpy as np
+import sounddevice as sd
+
+from livekit import rtc, api
+from livekit.plugins import noise_cancellation
+
+SAMPLERATE = 48000
+BLOCKSIZE = 480  # 10ms chunks at 48kHz
+CHANNELS = 1
+
+
+class AudioBuffer:
+    def __init__(self, blocksize=BLOCKSIZE):
+        self.blocksize = blocksize
+        self.buffer = np.array([], dtype=np.int16)
+
+    def add_frame(self, frame_data):
+        self.buffer = np.concatenate([self.buffer, frame_data])
+
+    def get_chunk(self):
+        if len(self.buffer) >= self.blocksize:
+            chunk = self.buffer[: self.blocksize]
+            self.buffer = self.buffer[self.blocksize :]
+            return chunk
+        return None
+
+    def get_padded_chunk(self):
+        if len(self.buffer) > 0:
+            chunk = np.zeros(self.blocksize, dtype=np.int16)
+            available = min(len(self.buffer), self.blocksize)
+            chunk[:available] = self.buffer[:available]
+            self.buffer = self.buffer[available:]
+            return chunk
+        return np.zeros(self.blocksize, dtype=np.int16)
+
+
+async def audio_player(queue: asyncio.Queue):
+    """Pull from the queue and stream audio using sounddevice."""
+    buffer = AudioBuffer(BLOCKSIZE)
+
+    def callback(outdata, frames, time, status):
+        if status:
+            print(f"Audio callback status: {status}")
+
+        # Try to fill buffer from queue
+        while not queue.empty():
+            try:
+                data = queue.get_nowait()
+                buffer.add_frame(data)
+            except asyncio.QueueEmpty:
+                break
+
+        # Get exactly the right amount of data
+        chunk = buffer.get_chunk()
+        if chunk is not None:
+            outdata[:] = chunk.reshape(-1, 1)
+        else:
+            # Not enough data, use what we have padded with zeros
+            outdata[:] = buffer.get_padded_chunk().reshape(-1, 1)
+
+    stream = sd.OutputStream(
+        samplerate=SAMPLERATE,
+        channels=CHANNELS,
+        blocksize=BLOCKSIZE,
+        dtype="int16",
+        callback=callback,
+        latency="low",
+    )
+    with stream:
+        while True:
+            await asyncio.sleep(0.1)  # keep the loop alive
+
+
+async def rtc_session(room, queue: asyncio.Queue):
+    track: rtc.RemoteAudioTrack | None = None
+    while not track:
+        for participant in room.remote_participants.values():
+            for t in participant.track_publications.values():
+                if t.kind == rtc.TrackKind.KIND_AUDIO and t.subscribed:
+                    track = t.track
+                    break
+            if track:
+                break
+        if not track:
+            print("waiting for audio track")
+            await asyncio.sleep(2)
+
+    stream = rtc.AudioStream.from_track(
+        track=track,
+        sample_rate=SAMPLERATE,
+        num_channels=1,
+        noise_cancellation=noise_cancellation.BVC(),  # or NC()
+    )
+
+    print("playing stream")
+    try:
+        # Process audio frames from the stream
+        async for audio_frame_event in stream:
+            frame = audio_frame_event.frame
+
+            audio_data = np.frombuffer(frame.data, dtype=np.int16)
+
+            try:
+                await queue.put(audio_data)
+            except asyncio.QueueFull:
+                # Skip this frame if queue is full
+                print("Warning: Audio queue full, dropping frame")
+                continue
+
+    finally:
+        # Clean up the stream when done
+        await stream.aclose()
+
+
+async def main():
+    queue = asyncio.Queue(maxsize=50)
+    player_task = asyncio.create_task(audio_player(queue))
+
+    token = (
+        api.AccessToken()
+        .with_identity("python-bot")
+        .with_name("Python Bot")
+        .with_grants(
+            api.VideoGrants(
+                room_join=True,
+                room="my-room",
+                agent=True,
+            )
+        )
+        .to_jwt()
+    )
+    url = os.getenv("LIVEKIT_URL")
+
+    room = rtc.Room()
+    await room.connect(
+        url,
+        token,
+        options=rtc.RoomOptions(
+            auto_subscribe=True,
+        ),
+    )
+    print(f"Connected to room: {room.name}")
+
+    try:
+        await rtc_session(room, queue)
+    finally:
+        # Clean up
+        await room.disconnect()
+        player_task.cancel()
+        try:
+            await player_task
+        except asyncio.CancelledError:
+            pass
+
+
+asyncio.run(main())