Add VP8 encoding through WEBRTC

pyproject.toml

@@ -37,6 +37,9 @@ redis = ">=4.3.5,<5.0.0"
 requests = "^2.33.1"
 typing-extensions = "^4.14.1"
 uvicorn = { version = "^0.34.0", extras = ["standard"] }
+aiortc = "^1.14.0"
+asyncio = "^4.0.0"
+av = "^16"
 
 
 [tool.poetry.group.dev.dependencies]

video_streamer/core/config.py

@@ -1,6 +1,6 @@
 import sys
 
-from typing import Any, Dict, Optional, Union, Tuple
+from typing import Any, Dict, List, Optional, Union, Tuple
 from pathlib import Path
 from pydantic import BaseModel, Field, validate_call, FilePath
 from pydantic_core import ValidationError
@@ -94,6 +94,11 @@ class SourceConfiguration(BaseModel):
         title="Authentication Configurations",
         default=AuthenticationConfiguration(type=None),
     )
+    allowed_origins: List[str] = Field(
+        title="Allowed Origins",
+        description="List of allowed origins for CORS",
+        default=[],
+    )
 
 class ServerConfiguration(BaseModel):
     """Server Configuration"""

video_streamer/core/streamer.py

@@ -3,6 +3,13 @@
 import queue
 import time
 from typing import Tuple, Generator, Any, Union
+from av import VideoFrame
+from aiortc import MediaStreamTrack
+import asyncio
+
+from fractions import Fraction
+
+import numpy as np
 
 from video_streamer.core.camera import TestCamera, LimaCamera, MJPEGCamera, VideoTestCamera, Camera, RedisCamera, EpicsPVACamera
 from video_streamer.core.config import SourceConfiguration
@@ -16,7 +23,7 @@ def __init__(self, config: SourceConfiguration, host: str, port: int, debug: boo
         self._debug = debug
         self._expt = 0.05
 
-    def start(self) -> Union[Generator, None, subprocess.Popen]:
+    def start(self) -> Union[Generator, None, subprocess.Popen, MediaStreamTrack]:
         pass
 
     def stop(self) -> None:
@@ -194,3 +201,72 @@ def stop(self) -> None:
                 self._poll_image_p.kill()
         else:
             print("Streamer stopped properly")
+
+
+class VideoMediaStreamer(Streamer):
+    """
+    A video track that returns frames from a multiprocessing queue.
+    """
+
+    kind = "video"
+
+    def __init__(self, config, host, port, debug):
+        super().__init__(config, host, port, debug)
+
+        self._poll_image_p = None
+
+    def start(self) -> MediaStreamTrack:
+        camera = self.get_camera()
+
+        _q = multiprocessing.Queue(1)
+        self._poll_image_p = multiprocessing.Process(
+            target=camera.poll_image, args=(_q,)
+        )
+        self._poll_image_p.start()
+
+        return QueueVideoTrack(_q, camera.size)
+
+    def stop(self) -> None:
+        print("Stopping Streamer...")
+        if self._poll_image_p:
+            self._poll_image_p.terminate()
+
+        time.sleep(1)
+        try:
+            if self._poll_image_p and self._poll_image_p.is_alive():
+                raise Exception("Image poll process did not stop properly")
+        except Exception:
+            if self._poll_image_p:
+                self._poll_image_p.kill()
+        else:
+            print("Streamer stopped properly")
+
+class QueueVideoTrack(MediaStreamTrack):
+    kind = "video"
+
+    def __init__(self, q: multiprocessing.Queue, image_size:Tuple , fps: int = 15, format_hint: str = "rgb24"):
+        super().__init__()
+        self.q = q
+        self.image_size = image_size
+        self.fps = fps
+        self.time_base = 1 / fps
+        self.format_hint = format_hint
+        self._ts = 0
+
+    async def recv(self) -> VideoFrame:
+        frame_bytes = bytearray()
+        try:
+            while True:
+                frame_bytes = self.q.get_nowait()
+        except Exception:
+            if len(frame_bytes) == 0:
+                frame_bytes = await asyncio.get_event_loop().run_in_executor(None, self.q.get)
+
+        width, height = self.image_size
+        arr = np.frombuffer(frame_bytes, dtype=np.uint8).reshape((height, width, 3))
+
+        vf = VideoFrame.from_ndarray(arr, format=self.format_hint)
+        self._ts += 1
+        vf.pts = self._ts
+        vf.time_base = Fraction(1, self.fps)
+        return vf

video_streamer/main.py

@@ -71,7 +71,7 @@ def parse_args() -> argparse.Namespace:
         "-of",
         "--output-format",
         dest="output_format",
-        help="output format, MPEG1 or MJPEG",
+        help="output format, MPEG1, MJPEG or VP8",
         default="MPEG1",
     )
 
@@ -157,6 +157,15 @@ def parse_args() -> argparse.Namespace:
         default="",
     )
 
+    opt_parser.add_argument(
+        "-ao",
+        "--allowed_origins",
+        nargs="+",
+        dest="allowed_origins",
+        help="Allowed origins for CORS",
+        default=[],
+    )
+
     return opt_parser.parse_args()
 
 
@@ -189,7 +198,8 @@ def run() -> None:
                         "type": args.auth_type,
                         "username": args.username,
                         "password": args.password
-                    })
+                    }),
+                    "allowed_origins": args.allowed_origins
                 }
             }
         }
@@ -202,7 +212,6 @@ def run() -> None:
 
     for uri, source_config in config.sources.items():
         host, port = uri.split(":")
-
         app = create_app(source_config, host, int(port), debug=args.debug)
 
         if app:

video_streamer/server.py

@@ -1,11 +1,12 @@
 import os
+import asyncio
 
 from fastapi import FastAPI, WebSocket, Request, WebSocketDisconnect
 from fastapi.responses import StreamingResponse, HTMLResponse
 from fastapi.staticfiles import StaticFiles
 
 from video_streamer.core.websockethandler import WebsocketHandler
-from video_streamer.core.streamer import FFMPGStreamer, MJPEGStreamer
+from video_streamer.core.streamer import FFMPGStreamer, MJPEGStreamer, VideoMediaStreamer
 from fastapi.templating import Jinja2Templates
 
 from contextlib import asynccontextmanager
@@ -15,6 +16,10 @@
 from typing import Optional
 from types import FrameType
 
+from aiortc import RTCPeerConnection, RTCRtpSender, RTCSessionDescription
+from fastapi.middleware.cors import CORSMiddleware
+from fastapi.responses import JSONResponse
+
 # This function makes sure that the server is correctly shutted down
 def handle_shutdown(signum: int, frame: Optional[FrameType]) -> None:
     print(f"Received signal {signum}, shutting down streamer...")
@@ -116,5 +121,72 @@ async def video_in(request: Request):
 
     return app
 
+def create_vp8_app(config, host, port, debug):
+
+    video_capabilities = RTCRtpSender.getCapabilities("video")
+
+    vp8_codecs = [
+        codec for codec in video_capabilities.codecs
+        if codec.mimeType == "video/VP8"
+    ]
+
+    @asynccontextmanager
+    async def lifespan(app: FastAPI):
+        app.state.streamer = VideoMediaStreamer(config, host, port, debug)
+        app.state.video_track = app.state.streamer.start()
+
+        app.state.pcs = set()
+
+        # Some signals are catched from uvicorn/fastapi, this makes sure that the server is correctly 
+        # shutting down and the second part of the lifespan is called
+        signal.signal(signal.SIGTERM, handle_shutdown)
+        signal.signal(signal.SIGINT, handle_shutdown)
+        try:
+            yield
+        finally:
+            for pc in set(app.state.pcs):
+                try:
+                    await asyncio.wait_for(pc.close(), timeout=2)
+                except asyncio.TimeoutError:
+                    print("Timeout while closing peer connection, closing forcefully...")
+            app.state.pcs.clear()
+            app.state.streamer.stop()
+
+    app = FastAPI(lifespan=lifespan)
+    app.add_middleware(CORSMiddleware, allow_methods=["OPTIONS", "POST"], allow_origins=config.allowed_origins, allow_headers=["Content-Type"])
+
+    @app.post("/offer")
+    async def offer(request: Request):
+        params = await request.json()
+        offer = RTCSessionDescription(sdp=params["sdp"], type=params["type"])
+        pc = RTCPeerConnection()
+        request.app.state.pcs.add(pc)
+
+        track = request.app.state.video_track
+        pc.addTrack(track)
+
+        # set the preferred codec to VP8
+        for transceiver in pc.getTransceivers():
+            if transceiver.kind == "video":
+                transceiver.setCodecPreferences(vp8_codecs)
+
+        @pc.on("connectionstatechange")
+        async def on_state_change():
+            if pc.connectionState in ("failed", "closed", "disconnected"):
+                await pc.close()
+                request.app.state.pcs.discard(pc)
+
+        await pc.setRemoteDescription(offer)
+        answer = await pc.createAnswer()
+        await pc.setLocalDescription(answer)
+
+        return JSONResponse(
+            content={
+                "sdp": pc.localDescription.sdp,
+                "type": pc.localDescription.type,
+            }
+        )
+
+    return app
 
-available_applications = {"MPEG1": create_mpeg1_app, "MJPEG": create_mjpeg_app}
+available_applications = {"MPEG1": create_mpeg1_app, "MJPEG": create_mjpeg_app, "VP8": create_vp8_app}