Add support for local stream processing to webrtc sdk (#2200)

* Add support for local stream processing to webrtc sdk

* fixed async blocking issue

* added timeout rtsp

* Add paced track for rtsp streams to address corrupted frames

---------

Co-authored-by: Rafel Bennasar Crespi <253519461+rafel-roboflow@users.noreply.github.com>

Author: grzegorz-roboflow

examples/webrtc_sdk/local_stream_basic.py

@@ -0,0 +1,108 @@
+"""
+Minimal sample using the SDK's WebRTC namespace to stream frames from a local
+RTSP or RTMP camera to an inference server with WebRTC worker enabled.
+
+Unlike RTSPSource (where the server captures the stream), LocalStreamSource
+captures the stream locally and sends frames to the server. Use this when:
+- The camera is only accessible from your local network
+- The server cannot reach the camera directly
+
+Supported protocols:
+- RTSP: rtsp://host/stream or rtsps://host/stream
+- RTMP: rtmp://host/stream or rtmps://host/stream
+
+Usage:
+  python examples/webrtc_sdk/local_stream_basic.py \\
+      --stream-url rtsp://camera.local/stream \\
+      --workspace-name <your_workspace> \\
+      --workflow-id <your_workflow_id> \\
+      [--api-url http://localhost:9001] \\
+      [--api-key <ROBOFLOW_API_KEY>] \\
+      [--stream-output <output_name>] \\
+      [--data-output <output_name>]
+
+Press 'q' in the preview window to exit.
+"""
+import argparse
+
+import av
+import cv2
+
+from inference_sdk import InferenceHTTPClient
+from inference_sdk.webrtc import LocalStreamSource, StreamConfig, VideoMetadata
+
+# Suppress FFmpeg warnings from PyAV
+av.logging.set_level(av.logging.ERROR)
+
+
+def parse_args() -> argparse.Namespace:
+    p = argparse.ArgumentParser("WebRTC SDK local_stream_basic")
+    p.add_argument(
+        "--stream-url",
+        required=True,
+        help="Stream URL (rtsp://, rtsps://, rtmp://, or rtmps://)",
+    )
+    p.add_argument("--api-url", default="https://serverless.roboflow.com")
+    p.add_argument("--workspace-name", required=True)
+    p.add_argument("--workflow-id", required=True)
+    p.add_argument("--image-input-name", default="image")
+    p.add_argument("--api-key", default=None)
+    p.add_argument(
+        "--stream-output",
+        default=None,
+        help="Name of the workflow output to stream (e.g., 'image_output')",
+    )
+    p.add_argument(
+        "--data-output",
+        default=None,
+        help="Name of the workflow output to receive via data channel",
+    )
+    return p.parse_args()
+
+
+def main() -> None:
+    args = parse_args()
+    client = InferenceHTTPClient.init(api_url=args.api_url, api_key=args.api_key)
+
+    # Prepare source - captures stream locally and sends to server
+    source = LocalStreamSource(args.stream_url)
+
+    # Prepare config
+    stream_output = [args.stream_output] if args.stream_output else []
+    data_output = [args.data_output] if args.data_output else []
+    config = StreamConfig(stream_output=stream_output, data_output=data_output)
+
+    # Create streaming session
+    session = client.webrtc.stream(
+        source=source,
+        workflow=args.workflow_id,
+        workspace=args.workspace_name,
+        image_input=args.image_input_name,
+        config=config,
+    )
+
+    # Register frame handler
+    @session.on_frame
+    def show_frame(frame, metadata):
+        cv2.imshow("WebRTC SDK - Local Stream", frame)
+        if cv2.waitKey(1) & 0xFF == ord("q"):
+            session.close()  # Close session and cleanup resources
+
+    # Register data handlers
+    # Global handler (receives entire serialized_output_data dict + metadata)
+    @session.on_data()
+    def on_message(data: dict, metadata: VideoMetadata):
+        print(f"Frame {metadata.frame_id}: {data}")
+
+    # Field-specific handler example (uncomment and customize based on your workflow):
+    # @session.on_data("predictions")
+    # def on_predictions(predictions: dict, metadata: VideoMetadata):
+    #     print(f"Frame {metadata.frame_id} predictions: {predictions}")
+
+    # Run the session (auto-starts, blocks until close() is called or stream ends)
+    # Automatically closes on exception or when stream ends
+    session.run()
+
+
+if __name__ == "__main__":
+    main()

inference_sdk/webrtc/__init__.py

@@ -4,6 +4,7 @@
 from .config import StreamConfig  # noqa: F401
 from .session import VideoMetadata, WebRTCSession  # noqa: F401
 from .sources import (  # noqa: F401
+    LocalStreamSource,
     ManualSource,
     MJPEGSource,
     RTSPSource,
@@ -22,6 +23,7 @@
     # Source classes
     "StreamSource",
     "WebcamSource",
+    "LocalStreamSource",
     "RTSPSource",
     "MJPEGSource",
     "VideoFileSource",

inference_sdk/webrtc/sources.py

@@ -5,12 +5,16 @@
 """
 
 import asyncio
+import threading
+import time
 from abc import ABC, abstractmethod
 from typing import TYPE_CHECKING, Any, Callable, Dict, Optional, Tuple
 
 import cv2
 import numpy as np
 from aiortc import RTCPeerConnection, VideoStreamTrack
+from aiortc.contrib.media import MediaPlayer
+from aiortc.mediastreams import VIDEO_CLOCK_RATE, VIDEO_PTIME, VIDEO_TIME_BASE
 from av import VideoFrame
 
 from inference_sdk.http.errors import InvalidParameterError
@@ -181,6 +185,137 @@ async def cleanup(self) -> None:
             self._track.release()
 
 
+class _PacedTrack(VideoStreamTrack):
+    """Wraps a source track and paces frame delivery like VideoStreamTrack.
+
+    MediaPlayer's PlayerStreamTrack does NOT pace RTSP frames (RTSP is in
+    REAL_TIME_FORMATS so _throttle_playback is False).  This means frames
+    are pulled from FFmpeg and forwarded to the RTP sender as fast as
+    possible, causing packet bursts that overflow the receiver's jitter
+    buffer.
+
+    This wrapper adds the same ~33 ms sleep between frames that the base
+    VideoStreamTrack uses, preventing bursts.
+    """
+
+    def __init__(self, source):  # noqa: ANN001
+        super().__init__()
+        self._source = source
+        self._start_time: float = 0.0
+        self._pts = 0
+
+    async def recv(self) -> VideoFrame:
+        from aiortc.mediastreams import MediaStreamError
+
+        if self.readyState != "live":
+            raise MediaStreamError
+
+        frame = await self._source.recv()
+
+        # Pace delivery: sleep until the next frame slot (~33 ms apart)
+        if self._pts == 0:
+            self._start_time = time.time()
+        else:
+            wait = self._start_time + (self._pts / VIDEO_CLOCK_RATE) - time.time()
+            if wait > 0:
+                await asyncio.sleep(wait)
+
+        frame.pts = self._pts
+        frame.time_base = VIDEO_TIME_BASE
+        self._pts += int(VIDEO_PTIME * VIDEO_CLOCK_RATE)
+        return frame
+
+    def stop(self) -> None:
+        super().stop()
+        self._source.stop()
+
+
+class LocalStreamSource(StreamSource):
+    """Stream source for locally captured RTSP/RTMP camera streams.
+
+    Unlike RTSPSource (where the server captures the RTSP stream), this source
+    captures the stream locally using aiortc's MediaPlayer (FFmpeg-based) and
+    sends frames to the server via WebRTC, similar to WebcamSource.
+
+    Supported protocols:
+    - RTSP: rtsp://host/path or rtsps://host/path
+    - RTMP: rtmp://host/path or rtmps://host/path
+
+    Use this when:
+    - The camera is only accessible from the client machine (e.g., local network)
+    - You want to preprocess frames before sending to the server
+    - The server cannot access the stream URL directly
+    """
+
+    # Supported URL schemes
+    SUPPORTED_SCHEMES = ("rtsp://", "rtsps://", "rtmp://", "rtmps://")
+
+    def __init__(self, url: str):
+        """Initialize local stream source.
+
+        Args:
+            url: Stream URL. Supported formats:
+                - RTSP: "rtsp://host/stream" or "rtsps://host/stream"
+                - RTMP: "rtmp://host/stream" or "rtmps://host/stream"
+                Credentials can be included: "rtsp://user:pass@host/stream"
+        """
+        if not url.startswith(self.SUPPORTED_SCHEMES):
+            raise InvalidParameterError(
+                f"Invalid stream URL: {url}. "
+                f"Must start with one of: {', '.join(self.SUPPORTED_SCHEMES)}"
+            )
+        self.url = url
+        self._player: Optional[MediaPlayer] = None
+
+    async def configure_peer_connection(self, pc: RTCPeerConnection) -> None:
+        """Create MediaPlayer for stream and add video track to peer connection."""
+
+        if self.url.startswith(("rtsp://", "rtsps://")):
+            self._player = await asyncio.to_thread(
+                MediaPlayer,
+                self.url,
+                format="rtsp",
+                options={
+                    "rtsp_transport": "tcp",
+                    "rtsp_flags": "prefer_tcp",
+                    "stimeout": "5000000",  # 5s RTSP socket timeout
+                    "timeout": "5000000",  # 5s TCP connection timeout
+                },
+            )
+        else:
+            self._player = await asyncio.to_thread(
+                MediaPlayer,
+                self.url,
+                options={
+                    "rw_timeout": "5000000",  # 5s socket timeout
+                },
+            )
+
+        if self._player.video is None:
+            raise RuntimeError(f"No video track available from stream: {self.url}")
+
+        # Wrap in a pacing track — MediaPlayer does not pace RTSP frames
+        # (RTSP is in REAL_TIME_FORMATS so _throttle_playback is False).
+        # Without pacing the RTP sender bursts all packets at once,
+        # overflowing the receiver's jitter buffer → truncated VP8 frames.
+        self._paced_track = _PacedTrack(self._player.video)
+        pc.addTrack(self._paced_track)
+
+    def get_initialization_params(self, config: "StreamConfig") -> Dict[str, Any]:
+        """Return empty params - stream is captured locally, not by server."""
+        return {}
+
+    async def cleanup(self) -> None:
+        """Stop the paced track and MediaPlayer."""
+        if hasattr(self, "_paced_track") and self._paced_track:
+            self._paced_track.stop()
+            self._paced_track = None
+        if self._player:
+            if self._player.video:
+                self._player.video.stop()
+            self._player = None
+
+
 class RTSPSource(StreamSource):
     """Stream source for RTSP camera streams.