File size: 6,658 Bytes
7847f44 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 | """
WebRTC Stream Manager for real-time video streaming using aiortc.
"""
import asyncio
import logging
from typing import Dict, Optional
import numpy as np
import cv2
from aiortc import RTCPeerConnection, RTCSessionDescription, VideoStreamTrack
from aiortc.contrib.media import MediaRelay
from av import VideoFrame
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
class VideoFrameTrack(VideoStreamTrack):
"""
Custom video track that streams frames pushed from the server.
"""
kind = "video"
def __init__(self, fps: int = 30):
super().__init__()
self.fps = fps
self.frame_queue: asyncio.Queue = asyncio.Queue(maxsize=60)
self._timestamp = 0
self._time_base = 1 / fps
self._running = True
async def recv(self) -> VideoFrame:
"""Receive the next frame to send."""
try:
# Wait for a frame with timeout
frame_data = await asyncio.wait_for(self.frame_queue.get(), timeout=1.0)
# Convert numpy array (BGR) to VideoFrame
if isinstance(frame_data, np.ndarray):
# Convert BGR to RGB
frame_rgb = cv2.cvtColor(frame_data, cv2.COLOR_BGR2RGB)
frame = VideoFrame.from_ndarray(frame_rgb, format="rgb24")
else:
frame = frame_data
# Set timing
frame.pts = self._timestamp
frame.time_base = self._time_base
self._timestamp += 1
return frame
except asyncio.TimeoutError:
# Return a black frame if no data available
black_frame = np.zeros((256, 256, 3), dtype=np.uint8)
frame = VideoFrame.from_ndarray(black_frame, format="rgb24")
frame.pts = self._timestamp
frame.time_base = self._time_base
self._timestamp += 1
return frame
async def push_frame(self, frame: np.ndarray):
"""Push a frame to be sent via WebRTC."""
try:
# Drop oldest frame if queue is full
if self.frame_queue.full():
try:
self.frame_queue.get_nowait()
except asyncio.QueueEmpty:
pass
await self.frame_queue.put(frame)
except Exception as e:
logger.error(f"Error pushing frame: {e}")
def stop(self):
"""Stop the track."""
self._running = False
super().stop()
class WebRTCConnection:
"""Manages a single WebRTC peer connection."""
def __init__(self, session_id: str, fps: int = 30):
self.session_id = session_id
self.pc = RTCPeerConnection()
self.video_track = VideoFrameTrack(fps=fps)
self.connected = False
self._connection_event = asyncio.Event()
# Add video track to peer connection
self.pc.addTrack(self.video_track)
# Handle connection state changes
@self.pc.on("connectionstatechange")
async def on_connectionstatechange():
state = self.pc.connectionState
logger.info(f"[WebRTC:{session_id}] Connection state: {state}")
if state == "connected":
self.connected = True
self._connection_event.set()
elif state in ("failed", "closed", "disconnected"):
self.connected = False
self._connection_event.clear()
@self.pc.on("iceconnectionstatechange")
async def on_iceconnectionstatechange():
logger.info(f"[WebRTC:{session_id}] ICE state: {self.pc.iceConnectionState}")
async def handle_offer(self, sdp: str, type: str) -> dict:
"""Handle an SDP offer and return an answer."""
offer = RTCSessionDescription(sdp=sdp, type=type)
await self.pc.setRemoteDescription(offer)
answer = await self.pc.createAnswer()
await self.pc.setLocalDescription(answer)
return {
"sdp": self.pc.localDescription.sdp,
"type": self.pc.localDescription.type
}
async def send_frame(self, frame: np.ndarray):
"""Send a video frame."""
if self.connected:
await self.video_track.push_frame(frame)
async def wait_for_connection(self, timeout: float = 5.0) -> bool:
"""Wait for the connection to be established."""
try:
await asyncio.wait_for(self._connection_event.wait(), timeout=timeout)
return True
except asyncio.TimeoutError:
return False
async def close(self):
"""Close the connection."""
self.video_track.stop()
await self.pc.close()
self.connected = False
class WebRTCManager:
"""Manages multiple WebRTC connections."""
def __init__(self):
self.connections: Dict[str, WebRTCConnection] = {}
self.relay = MediaRelay()
async def handle_offer(self, session_id: str, sdp: str, type: str, fps: int = 30) -> dict:
"""Handle a WebRTC offer for a session."""
# Close existing connection if any
if session_id in self.connections:
await self.close_connection(session_id)
# Create new connection
connection = WebRTCConnection(session_id, fps=fps)
self.connections[session_id] = connection
# Handle the offer
answer = await connection.handle_offer(sdp, type)
logger.info(f"[WebRTC] Created connection for session {session_id}")
return answer
def is_connected(self, session_id: str) -> bool:
"""Check if a session is connected."""
connection = self.connections.get(session_id)
return connection.connected if connection else False
async def wait_for_connection(self, session_id: str, timeout: float = 5.0) -> bool:
"""Wait for a session to be connected."""
connection = self.connections.get(session_id)
if not connection:
return False
return await connection.wait_for_connection(timeout)
async def send_frame(self, session_id: str, frame: np.ndarray):
"""Send a frame to a session."""
connection = self.connections.get(session_id)
if connection and connection.connected:
await connection.send_frame(frame)
async def close_connection(self, session_id: str):
"""Close a session's connection."""
connection = self.connections.pop(session_id, None)
if connection:
await connection.close()
logger.info(f"[WebRTC] Closed connection for session {session_id}")
# Global manager instance
webrtc_manager = WebRTCManager()
|