Programmable video source (whitphx#1349)

whitphx · web-flow · commit 14992e2841c5 · 2023-08-21T18:03:14.000+09:00
* Create VideoSourceTrack

* Remove queue

* Change the frame buffer to be an instance of av.VideoFrame

* Fix VideoSourceTrack to be callback-based

* Fix the factory function

* Fix

* Stop the source tracks automatically

* Stop the relayed source tracks automatically

* Stop the source track in the sample

* Pass pts and time_base to the callback

* Make fps configurable

* Fix

* Export VideoSourceCallback

* Fix

* Fix type annotations

* Fix

* Fix
diff --git a/pages/14_programmable_source.py b/pages/14_programmable_source.py
@@ -0,0 +1,61 @@
+import fractions
+import time
+
+import av
+import cv2
+import numpy as np
+import streamlit as st
+from streamlit_webrtc import WebRtcMode, create_video_source_track, webrtc_streamer
+
+thickness = st.slider("thickness", 1, 10, 3, 1)
+
+
+def video_source_callback(pts: int, time_base: fractions.Fraction) -> av.VideoFrame:
+    pts_sec = pts * time_base
+
+    buffer = np.zeros((480, 640, 3), dtype=np.uint8)
+    buffer = cv2.putText(
+        buffer,
+        text=f"time: {time.time():.2f}",
+        org=(0, 32),
+        fontFace=cv2.FONT_HERSHEY_SIMPLEX,
+        fontScale=1.0,
+        color=(255, 255, 0),
+        thickness=thickness,
+        lineType=cv2.LINE_4,
+    )
+    buffer = cv2.putText(
+        buffer,
+        text=f"pts: {pts} ({float(pts_sec):.2f} sec)",
+        org=(0, 64),
+        fontFace=cv2.FONT_HERSHEY_SIMPLEX,
+        fontScale=1.0,
+        color=(255, 255, 0),
+        thickness=thickness,
+        lineType=cv2.LINE_4,
+    )
+    return av.VideoFrame.from_ndarray(buffer, format="bgr24")
+
+
+fps = st.slider("fps", 1, 30, 30, 1)
+
+
+video_source_track = create_video_source_track(
+    video_source_callback, key="video_source_track", fps=fps
+)
+
+
+def on_change():
+    ctx = st.session_state["player"]
+    stopped = not ctx.state.playing and not ctx.state.signalling
+    if stopped:
+        video_source_track.stop()  # Manually stop the track.
+
+
+webrtc_streamer(
+    key="player",
+    mode=WebRtcMode.RECVONLY,
+    source_video_track=video_source_track,
+    media_stream_constraints={"video": True, "audio": False},
+    on_change=on_change,
+)
diff --git a/streamlit_webrtc/__init__.py b/streamlit_webrtc/__init__.py
@@ -23,8 +23,9 @@
     Translations,
     VideoHTMLAttributes,
 )
-from .factory import create_mix_track, create_process_track
+from .factory import create_mix_track, create_process_track, create_video_source_track
 from .mix import MixerCallback
+from .source import VideoSourceCallback, VideoSourceTrack
 from .webrtc import (
     AudioProcessorBase,
     AudioProcessorFactory,
@@ -64,6 +65,9 @@
     "VideoProcessorFactory",
     "VideoTransformerBase",  # XXX: Deprecated
     "VideoReceiver",
+    "VideoSourceTrack",
+    "VideoSourceCallback",
+    "create_video_source_track",
     "WebRtcMode",
     "WebRtcWorker",
     "MediaStreamConstraints",
diff --git a/streamlit_webrtc/factory.py b/streamlit_webrtc/factory.py
@@ -31,6 +31,7 @@
     VideoProcessTrack,
 )
 from .relay import get_global_relay
+from .source import VideoSourceCallback, VideoSourceTrack
 
 _PROCESSOR_TRACK_CACHE_KEY_PREFIX = "__PROCESSOR_TRACK_CACHE__"
 
@@ -191,3 +192,27 @@ def create_mix_track(
         )
         st.session_state[cache_key] = mixer_track
     return mixer_track
+
+
+_VIDEO_SOURCE_TRACK_CACHE_KEY_PREFIX = "__VIDEO_SOURCE_TRACK_CACHE__"
+
+
+def create_video_source_track(
+    callback: VideoSourceCallback,
+    key: str,
+    fps=30,
+) -> VideoSourceTrack:
+    cache_key = _VIDEO_SOURCE_TRACK_CACHE_KEY_PREFIX + key
+    if (
+        cache_key in st.session_state
+        and isinstance(st.session_state[cache_key], VideoSourceTrack)
+        and st.session_state[cache_key].kind == "video"
+        and st.session_state[cache_key].readyState == "live"
+    ):
+        video_source_track: VideoSourceTrack = st.session_state[cache_key]
+        video_source_track._callback = callback
+        video_source_track._fps = fps
+    else:
+        video_source_track = VideoSourceTrack(callback=callback, fps=fps)
+        st.session_state[cache_key] = video_source_track
+    return video_source_track
diff --git a/streamlit_webrtc/source.py b/streamlit_webrtc/source.py
@@ -0,0 +1,73 @@
+import asyncio
+import fractions
+import logging
+import time
+from typing import Callable, Optional, Union
+
+import av
+from aiortc import MediaStreamTrack
+from aiortc.mediastreams import MediaStreamError
+
+logger = logging.getLogger(__name__)
+
+# Copied from https://github.com/aiortc/aiortc/blob/main/src/aiortc/mediastreams.py
+AUDIO_PTIME = 0.020  # 20ms audio packetization
+VIDEO_CLOCK_RATE = 90000
+VIDEO_TIME_BASE = fractions.Fraction(1, VIDEO_CLOCK_RATE)
+
+
+# Ref: VideoStreamTrack and AudioStreamTrack in
+# https://github.com/aiortc/aiortc/blob/main/src/aiortc/mediastreams.py
+
+
+VideoSourceCallback = Callable[
+    [int, fractions.Fraction], av.VideoFrame
+]  # (pts, time_base) -> frame
+
+
+class VideoSourceTrack(MediaStreamTrack):
+    def __init__(self, callback: VideoSourceCallback, fps: Union[int, float]) -> None:
+        super().__init__()
+        self.kind = "video"
+        self._callback = callback
+        self._fps = fps
+        self._started_at: Optional[float] = None
+        self._pts: Optional[int] = None
+
+    async def recv(self) -> av.frame.Frame:
+        if self.readyState != "live":
+            raise MediaStreamError
+
+        if self._started_at is None or self._pts is None:
+            self._started_at = time.monotonic()
+            self._pts = 0
+
+            frame = self._call_callback(self._pts, VIDEO_TIME_BASE)
+        else:
+            self._pts += int(VIDEO_CLOCK_RATE / self._fps)
+
+            frame = self._call_callback(self._pts, VIDEO_TIME_BASE)
+
+            wait = self._started_at + (self._pts / VIDEO_CLOCK_RATE) - time.monotonic()
+            if wait < 0:
+                logger.warning(
+                    "VideoSourceCallbackTrack: Video frame callback is too slow."
+                )
+                wait = 0
+            await asyncio.sleep(wait)
+
+        return frame
+
+    def _call_callback(self, pts: int, time_base: fractions.Fraction) -> av.VideoFrame:
+        try:
+            frame = self._callback(pts, time_base)
+        except Exception as exc:
+            logger.error(
+                "VideoSourceCallbackTrack: Video frame callback raised an exception: %s",  # noqa: E501
+                exc,
+            )
+            raise
+
+        frame.pts = pts
+        frame.time_base = time_base
+        return frame
diff --git a/streamlit_webrtc/webrtc.py b/streamlit_webrtc/webrtc.py
@@ -308,6 +308,8 @@ class WebRtcWorker(Generic[VideoProcessorT, AudioProcessorT]):
     _output_video_track: Optional[MediaStreamTrack]
     _output_audio_track: Optional[MediaStreamTrack]
     _player: Optional[MediaPlayer]
+    _relayed_source_video_track: Optional[MediaRelay]
+    _relayed_source_audio_track: Optional[MediaRelay]
 
     @property
     def video_processor(
@@ -404,6 +406,8 @@ def __init__(
         self._output_video_track = None
         self._output_audio_track = None
         self._player = None
+        self._relayed_source_video_track = None
+        self._relayed_source_audio_track = None
 
         self._session_shutdown_observer = SessionShutdownObserver(self.stop)
 
@@ -504,14 +508,20 @@ def on_track_created(track_type: TrackType, track: MediaStreamTrack):
             player = self.player_factory()
             self._player = player
             if player.audio:
-                source_audio_track = relay.subscribe(player.audio)
+                source_audio_track = player.audio
             if player.video:
-                source_video_track = relay.subscribe(player.video)
+                source_video_track = player.video
         else:
-            if self.source_video_track:
-                source_video_track = relay.subscribe(self.source_video_track)
             if self.source_audio_track:
-                source_audio_track = relay.subscribe(self.source_audio_track)
+                self._relayed_source_audio_track = relay.subscribe(
+                    self.source_audio_track
+                )
+                source_audio_track = self._relayed_source_audio_track
+            if self.source_video_track:
+                self._relayed_source_video_track = relay.subscribe(
+                    self.source_video_track
+                )
+                source_video_track = self._relayed_source_video_track
 
         @self.pc.on("iceconnectionstatechange")
         async def on_iceconnectionstatechange():
@@ -655,6 +665,21 @@ def _unset_processors(self):
                 self._player.audio.stop()
         self._player = None
 
+        # Same as above,
+        # the source tracks are not automatically stopped when the WebRTC.
+        # Only the relayed tracks are stopped here because
+        # the upstream tracks may still be used by other consumers.
+        if self._relayed_source_audio_track:
+            logger.debug("Stopping the relayed source audio track")
+            self._relayed_source_audio_track.stop()
+        self.source_audio_track = None
+        self._relayed_source_audio_track = None
+        if self._relayed_source_video_track:
+            logger.debug("Stopping the relayed source video track")
+            self._relayed_source_video_track.stop()
+        self.source_video_track = None
+        self._relayed_source_video_track = None
+
     def stop(self, timeout: Union[float, None] = 1.0):
         self._unset_processors()
         if self._process_offer_thread:
