openpilot0/tools/streamer/streamerd.py

90 lines
2.5 KiB
Python
Executable File

#!/usr/bin/env python
# pylint: skip-file
import os
import sys
import zmq
import cv2
import numpy as np
import struct
# sudo pip install git+git://github.com/mikeboers/PyAV.git
import av
import cereal.messaging as messaging
from cereal.services import service_list
PYGAME = os.getenv("PYGAME") is not None
if PYGAME:
import pygame
imgff = np.zeros((874, 1164, 3), dtype=np.uint8)
# first 74 bytes in any stream
start = "0000000140010c01ffff016000000300b0000003000003005dac5900000001420101016000000300b0000003000003005da0025080381c5c665aee4c92ec80000000014401c0f1800420"
def receiver_thread():
if PYGAME:
pygame.init()
pygame.display.set_caption("vnet debug UI")
screen = pygame.display.set_mode((1164, 874), pygame.DOUBLEBUF)
camera_surface = pygame.surface.Surface((1164, 874), 0, 24).convert()
addr = "192.168.5.11"
if len(sys.argv) >= 2:
addr = sys.argv[1]
context = zmq.Context()
s = messaging.sub_sock(context, 9002, addr=addr)
frame_sock = messaging.pub_sock(context, service_list['frame'].port)
ctx = av.codec.codec.Codec('hevc', 'r').create()
ctx.decode(av.packet.Packet(start.decode("hex")))
# import time
while 1:
# t1 = time.time()
ts, raw = s.recv_multipart()
ts = struct.unpack('q', ts)[0] * 1000
# t1, t2 = time.time(), t1
#print 'ms to get frame:', (t1-t2)*1000
pkt = av.packet.Packet(raw)
f = ctx.decode(pkt)
if not f:
continue
f = f[0]
# t1, t2 = time.time(), t1
#print 'ms to decode:', (t1-t2)*1000
y_plane = np.frombuffer(f.planes[0], np.uint8).reshape((874, 1216))[:, 0:1164]
u_plane = np.frombuffer(f.planes[1], np.uint8).reshape((437, 608))[:, 0:582]
v_plane = np.frombuffer(f.planes[2], np.uint8).reshape((437, 608))[:, 0:582]
yuv_img = y_plane.tobytes() + u_plane.tobytes() + v_plane.tobytes()
# t1, t2 = time.time(), t1
#print 'ms to make yuv:', (t1-t2)*1000
#print 'tsEof:', ts
dat = messaging.new_message('frame')
dat.frame.image = yuv_img
dat.frame.timestampEof = ts
dat.frame.transform = map(float, list(np.eye(3).flatten()))
frame_sock.send(dat.to_bytes())
if PYGAME:
yuv_np = np.frombuffer(yuv_img, dtype=np.uint8).reshape(874 * 3 // 2, -1)
cv2.cvtColor(yuv_np, cv2.COLOR_YUV2RGB_I420, dst=imgff)
#print yuv_np.shape, imgff.shape
#scipy.misc.imsave("tmp.png", imgff)
pygame.surfarray.blit_array(camera_surface, imgff.swapaxes(0, 1))
screen.blit(camera_surface, (0, 0))
pygame.display.flip()
def main(gctx=None):
receiver_thread()
if __name__ == "__main__":
main()