mirror of
https://github.com/comma-hacks/webrtc.git
synced 2025-10-01 06:22:10 +08:00
85 lines
2.4 KiB
Python
Executable File
85 lines
2.4 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
import cereal.messaging as messaging
|
|
from aiortc import VideoStreamTrack
|
|
import asyncio
|
|
import numpy as np
|
|
import av
|
|
import os
|
|
import sys
|
|
import time
|
|
|
|
W, H = 1928, 1208
|
|
V4L2_BUF_FLAG_KEYFRAME = 8
|
|
|
|
class VisionIpcTrack(VideoStreamTrack):
|
|
def __init__(self, sock_name, addr):
|
|
super().__init__()
|
|
self.codec = av.CodecContext.create("hevc", "r")
|
|
os.environ["ZMQ"] = "1"
|
|
messaging.context = messaging.Context()
|
|
self.sock = messaging.sub_sock(sock_name, None, addr=addr, conflate=False)
|
|
self.cnt = 0
|
|
self.last_idx = -1
|
|
self.seen_iframe = False
|
|
self.time_q = []
|
|
self.sock_name = sock_name
|
|
|
|
async def recv(self):
|
|
pts, time_base = await self.next_timestamp()
|
|
frame = None
|
|
while frame is None:
|
|
msgs = messaging.drain_sock(self.sock, wait_for_one=True)
|
|
for evt_context in msgs:
|
|
with evt_context as evt:
|
|
evta = getattr(evt, evt.which())
|
|
if evta.idx.encodeId != 0 and evta.idx.encodeId != (self.last_idx+1):
|
|
print("DROP PACKET!")
|
|
self.last_idx = evta.idx.encodeId
|
|
if not self.seen_iframe and not (evta.idx.flags & V4L2_BUF_FLAG_KEYFRAME):
|
|
print("waiting for iframe")
|
|
continue
|
|
self.time_q.append(time.monotonic())
|
|
|
|
# put in header (first)
|
|
if not self.seen_iframe:
|
|
self.codec.decode(av.packet.Packet(evta.header))
|
|
self.seen_iframe = True
|
|
|
|
frames = self.codec.decode(av.packet.Packet(evta.data))
|
|
if len(frames) == 0:
|
|
print("DROP SURFACE")
|
|
continue
|
|
assert len(frames) == 1
|
|
|
|
frame = frames[0]
|
|
frame.pts = pts
|
|
frame.time_base = time_base
|
|
return frame
|
|
|
|
if __name__ == "__main__":
|
|
from time import time_ns
|
|
import sys
|
|
|
|
async def test():
|
|
frame_count=0
|
|
start_time=time_ns()
|
|
track = VisionIpcTrack("roadEncodeData", "tici")
|
|
while True:
|
|
await track.recv()
|
|
now = time_ns()
|
|
playtime = now - start_time
|
|
playtime_sec = playtime * 0.000000001
|
|
if playtime_sec >= 1:
|
|
print(f'fps: {frame_count}')
|
|
frame_count = 0
|
|
start_time = time_ns()
|
|
else:
|
|
frame_count+=1
|
|
|
|
# Run event loop
|
|
loop = asyncio.new_event_loop()
|
|
try:
|
|
loop.run_until_complete(test())
|
|
except KeyboardInterrupt:
|
|
sys.exit(0)
|