Files
VisDrone-Version/rtc_handler.py
2025-08-05 16:55:45 +08:00

95 lines
2.9 KiB
Python

import asyncio
import queue
from fractions import Fraction
from urllib.parse import urlparse
import aiohttp
import av
import numpy as np
from aiortc import RTCPeerConnection, RTCSessionDescription, RTCConfiguration, VideoStreamTrack
class DummyVideoTrack(VideoStreamTrack):
async def recv(self):
# 简洁初始化、返回固定颜色的帧
return np.full((480, 640, 3), (0, 0, 255), dtype=np.uint8)
async def receive_video_frames(whep_url):
pc = RTCPeerConnection(RTCConfiguration(iceServers=[]))
frames_queue = asyncio.Queue()
pc.addTrack(DummyVideoTrack())
@pc.on("track")
def on_track(track):
if track.kind == "video":
asyncio.create_task(consume_track(track, frames_queue))
@pc.on("iceconnectionstatechange")
def on_ice_connection_state_change():
print(f"ICE 连接状态: {pc.iceConnectionState}")
offer = await pc.createOffer()
await pc.setLocalDescription(offer)
headers = {"Content-Type": "application/sdp"}
async with aiohttp.ClientSession() as session:
async with session.post(whep_url, data=pc.localDescription.sdp, headers=headers) as response:
if response.status != 201:
raise Exception(f"服务器返回错误: {response.status}")
answer = RTCSessionDescription(sdp=await response.text(), type="answer")
await pc.setRemoteDescription(answer)
if "Location" in response.headers:
base_url = f"{urlparse(whep_url).scheme}://{urlparse(whep_url).netloc}"
print("ICE 协商 URL:", base_url + response.headers["Location"])
while pc.iceConnectionState not in ["connected", "completed"]:
await asyncio.sleep(1)
print("ICE 连接完成,开始接收视频流")
try:
while True:
frame = await frames_queue.get()
if frame is None:
break
yield frame
except KeyboardInterrupt:
pass
finally:
await pc.close()
async def consume_track(track, frames_queue):
try:
while True:
frame = await track.recv()
if frame is None:
print("没有接收到有效的帧数据")
await frames_queue.put(None)
break
img = frame.to_ndarray(format="bgr24")
await frames_queue.put(img)
except Exception as e:
print("处理帧错误:", e)
await frames_queue.put(None)
def rtc_frame(url, frame_queue):
async def main():
async for frame in receive_video_frames(url):
try:
frame_queue.put_nowait(frame)
except queue.Full:
frame_queue.get_nowait()
frame_queue.put_nowait(frame)
loop = asyncio.new_event_loop()
asyncio.set_event_loop(loop)
loop.run_until_complete(main())
loop.close()