flandre/src/rtcserver.py

104 lines
2.9 KiB
Python
Raw Normal View History

2025-01-06 11:21:04 +08:00
import asyncio
import json
import logging
import os
from fractions import Fraction
import aiohttp_cors
import aiortc.rtcrtpsender
import zmq
from aiohttp import web
from aiortc import MediaStreamTrack, RTCPeerConnection, RTCSessionDescription, RTCConfiguration, RTCRtpCodecCapability
from av import VideoFrame
import H264NV
ROOT = os.path.dirname(__file__)
class PlayerStreamTrackx(MediaStreamTrack):
def __init__(self):
super().__init__()
context = zmq.Context()
self.socket = context.socket(zmq.PULL)
self.socket.bind("tcp://*:5555")
self.kind = 'video'
self.pts = 0
async def recv(self):
frame = VideoFrame.from_bytes(self.socket.recv(), 1920, 1080)
frame.time_base = Fraction(1, 120)
frame.pts = self.pts
self.pts += 1
return frame
pcs = set()
px = PlayerStreamTrackx()
async def offer(request):
params = await request.json()
offer = RTCSessionDescription(sdp=params["sdp"], type=params["type"])
pc = RTCPeerConnection(RTCConfiguration([]))
# print(request.remote)
pcs.add(pc)
rc = pc.addTransceiver(px, 'sendonly')
rc.setCodecPreferences([RTCRtpCodecCapability(mimeType='video/H264',
clockRate=90000,
channels=None,
parameters={
'level-asymmetry-allowed': '1',
'packetization-mode': '1',
'profile-level-id': '42e01f'
})])
await pc.setRemoteDescription(offer)
answer = await pc.createAnswer()
await pc.setLocalDescription(answer)
return web.Response(
content_type="application/json",
text=json.dumps(
{"sdp": pc.localDescription.sdp, "type": pc.localDescription.type}
),
)
async def on_shutdown(app):
# close peer connections
coros = [pc.close() for pc in pcs]
await asyncio.gather(*coros)
pcs.clear()
if __name__ == "__main__":
f0 = aiortc.RTCRtpSender.__init__
def f1(*args, **kwargs):
f0(*args, **kwargs)
if not args[1] == 'audio':
args[0]._RTCRtpSender__encoder = H264NV.H264NVENCEncoder()
aiortc.RTCRtpSender.__init__ = f1
logging.basicConfig(level=logging.INFO)
app = web.Application()
app.on_shutdown.append(on_shutdown)
app.router.add_post("/offer", offer)
cors = aiohttp_cors.setup(app, defaults={
"*": aiohttp_cors.ResourceOptions(
allow_credentials=True,
expose_headers="*",
allow_headers="*"
)
})
for route in list(app.router.routes()):
cors.add(route)
web.run_app(
app, access_log=None, host='0.0.0.0', port=8081
)