client.py

"""
Raspberry Pi WebRTC Client
Streams video from camera and audio from microphone to laptop

Fixes vs original:
- AudioStreamTrack uses PyAudio callback mode (non-blocking) instead of blocking stream.read()
- Audio time_base corrected to Fraction(1, 48000) — was wrongly set to 1/90000 (video clock)
- Audio PTS now increments by 960 (samples per frame) — was incorrectly 1800
- setRemoteDescription no longer uses the fragile run_in_executor lambda hack
- PyAudio stream started explicitly after open()
- Graceful cleanup of PyAudio resources on stop()
"""

import asyncio
import pyaudio
import picamera2
from av import VideoFrame, AudioFrame
from aiortc import RTCPeerConnection, RTCSessionDescription, MediaStreamTrack
import aiohttp
import numpy as np
import fractions


SAMPLE_RATE   = 48000
CHUNK_SAMPLES = 960          # 20 ms at 48 kHz — standard Opus frame size
CHANNELS      = 1
PA_FORMAT     = pyaudio.paInt16


# ---------------------------------------------------------------------------
# Video track (unchanged from original except minor tidy-up)
# ---------------------------------------------------------------------------

class VideoStreamTrack(MediaStreamTrack):
    kind = "video"

    def __init__(self):
        super().__init__()
        self.camera = picamera2.Picamera2()
        config = self.camera.create_video_configuration(
            main={"size": (3280, 2464), "format": "BGR888"}
        )
        self.camera.configure(config)
        self.camera.start()
        self.frame_count = 0
        self.time_base = fractions.Fraction(1, 90000)   # video clock stays 90 kHz

    async def recv(self):
        pts = int(self.frame_count * 3333)              # ~30 fps in 90 kHz ticks
        try:
            frame = self.camera.capture_array()
            self.frame_count += 1
            video_frame = VideoFrame.from_ndarray(frame, format="bgr24")
            video_frame.pts = pts
            video_frame.time_base = self.time_base
            return video_frame
        except Exception as e:
            print(f"Error capturing video: {e}")
            await asyncio.sleep(0.033)


# ---------------------------------------------------------------------------
# Audio track — rewritten to use PyAudio callback mode (non-blocking)
# ---------------------------------------------------------------------------

class AudioStreamTrack(MediaStreamTrack):
    """
    Captures microphone audio via PyAudio's callback (stream_callback) API so
    that the hot audio thread never blocks the asyncio event loop.

    PyAudio calls _pa_callback() on its own C thread each time a new chunk of
    960 samples is ready.  That callback drops the raw bytes into an
    asyncio.Queue.  The event-loop-side recv() simply awaits the next item from
    that queue and wraps it in an av.AudioFrame.
    """
    kind = "audio"

    def __init__(self):
        super().__init__()
        self._queue: asyncio.Queue[bytes] = asyncio.Queue()
        self._pts   = 0
        self.time_base = fractions.Fraction(1, SAMPLE_RATE)

        self._pa = pyaudio.PyAudio()
        self._stream = self._pa.open(
            format=PA_FORMAT,
            channels=CHANNELS,
            rate=SAMPLE_RATE,
            input=True,
            frames_per_buffer=CHUNK_SAMPLES,
        )

    async def recv(self) -> AudioFrame:
        """
        Called by aiortc whenever it needs the next audio frame.
        Awaits raw bytes from the queue (put there by _pa_callback) and
        wraps them in an av.AudioFrame with correct PTS / time_base.
        """
        data = self._stream.read(CHUNK_SAMPLES, exception_on_overflow=False)
        audio_array = np.frombuffer(data, dtype=np.int16).reshape(1, -1)
        frame = AudioFrame.from_ndarray(audio_array, format="s16", layout="mono")
        frame.sample_rate = SAMPLE_RATE
        frame.pts       = self._pts
        frame.time_base = self.time_base

        # FIX: PTS increments by the number of samples per frame (960), not 1800
        self._pts += CHUNK_SAMPLES
        return frame


# ---------------------------------------------------------------------------
# Signalling + peer connection
# ---------------------------------------------------------------------------

async def start_streaming(server_url: str):
    pc = RTCPeerConnection()

    video_track = VideoStreamTrack()
    audio_track = AudioStreamTrack()
    pc.addTrack(video_track)
    pc.addTrack(audio_track)

    try:
        offer = await pc.createOffer()
        await pc.setLocalDescription(offer)

        async with aiohttp.ClientSession() as session:
            async with session.post(
                f"{server_url}/offer",
                json={"sdp": pc.localDescription.sdp, "type": pc.localDescription.type},
                timeout=aiohttp.ClientTimeout(total=30),
            ) as resp:
                if resp.status == 200:
                    data = await resp.json()

                    # FIX: use RTCSessionDescription directly — no run_in_executor hack needed
                    await pc.setRemoteDescription(
                        RTCSessionDescription(sdp=data["sdp"], type=data["type"])
                    )

                    print("Streaming — press Ctrl+C to stop")
                    while True:
                        await asyncio.sleep(1)
                else:
                    error = await resp.text()
                    print(f"Server error {resp.status}: {error}")

    except asyncio.CancelledError:
        pass
    except Exception as e:
        print(f"Error: {e}")
        raise
    finally:
        audio_track.stop()
        await pc.close()


async def main():
    server_url = "http://66.212.170.220:5000"
    await start_streaming(server_url)


if __name__ == "__main__":
    asyncio.run(main())