compressed_vipc: make yuv default

old-commit-hash: 3694f9801c
taco
George Hotz 3 years ago
parent 8e755ff447
commit d348193169
  1. 39
      tools/camerastream/compressed_vipc.py

@ -7,6 +7,8 @@ import multiprocessing
from cereal.visionipc.visionipc_pyx import VisionIpcServer, VisionStreamType # pylint: disable=no-name-in-module, import-error from cereal.visionipc.visionipc_pyx import VisionIpcServer, VisionStreamType # pylint: disable=no-name-in-module, import-error
W, H = 1928, 1208 W, H = 1928, 1208
V4L2_BUF_FLAG_KEYFRAME = 8
def writer(fn, addr, sock_name): def writer(fn, addr, sock_name):
import cereal.messaging as messaging import cereal.messaging as messaging
HEADER = b"\x00\x00\x00\x01\x40\x01\x0c\x01\xff\xff\x01\x60\x00\x00\x03\x00\xb0\x00\x00\x03\x00\x00\x03\x00\x96\xac\x09\x00\x00\x00\x01\x42\x01\x01\x01\x60\x00\x00\x03\x00\xb0\x00\x00\x03\x00\x00\x03\x00\x96\xa0\x03\xd0\x80\x13\x07\x1b\x2e\x5a\xee\x4c\x92\xea\x00\xbb\x42\x84\xa0\x00\x00\x00\x01\x44\x01\xc0\xe2\x4f\x09\xc1\x80\xc6\x08\x40\x00" HEADER = b"\x00\x00\x00\x01\x40\x01\x0c\x01\xff\xff\x01\x60\x00\x00\x03\x00\xb0\x00\x00\x03\x00\x00\x03\x00\x96\xac\x09\x00\x00\x00\x01\x42\x01\x01\x01\x60\x00\x00\x03\x00\xb0\x00\x00\x03\x00\x00\x03\x00\x96\xa0\x03\xd0\x80\x13\x07\x1b\x2e\x5a\xee\x4c\x92\xea\x00\xbb\x42\x84\xa0\x00\x00\x00\x01\x44\x01\xc0\xe2\x4f\x09\xc1\x80\xc6\x08\x40\x00"
@ -29,7 +31,7 @@ def writer(fn, addr, sock_name):
if evta.idx != 0 and evta.idx != (last_idx+1): if evta.idx != 0 and evta.idx != (last_idx+1):
print("DROP!") print("DROP!")
last_idx = evta.idx last_idx = evta.idx
if len(evta.data) > 4 and evta.data[4] == 0x26: if evta.flags & V4L2_BUF_FLAG_KEYFRAME or evta.flags == 0x7f001030:
seen_iframe = True seen_iframe = True
if not seen_iframe: if not seen_iframe:
print("waiting for iframe") print("waiting for iframe")
@ -37,15 +39,16 @@ def writer(fn, addr, sock_name):
fifo_file.write(evta.data) fifo_file.write(evta.data)
fifo_file.flush() fifo_file.flush()
def decoder_nvidia(fn, vipc_server, vst, yuv=False): def decoder_nvidia(fn, vipc_server, vst, yuv=True, rgb=False):
sys.path.append("/raid.dell2/PyNvCodec") sys.path.append("/raid.dell2/PyNvCodec")
import PyNvCodec as nvc # pylint: disable=import-error import PyNvCodec as nvc # pylint: disable=import-error
decoder = nvc.PyNvDecoder(fn, 0, {"probesize": "32"}) decoder = nvc.PyNvDecoder(fn, 0, {"probesize": "32"})
cc1 = nvc.ColorspaceConversionContext(nvc.ColorSpace.BT_709, nvc.ColorRange.JPEG) cc1 = nvc.ColorspaceConversionContext(nvc.ColorSpace.BT_709, nvc.ColorRange.JPEG)
conv = nvc.PySurfaceConverter(W, H, nvc.PixelFormat.NV12, nvc.PixelFormat.BGR, 0) if rgb:
nvDwn = nvc.PySurfaceDownloader(W, H, nvc.PixelFormat.BGR, 0) conv = nvc.PySurfaceConverter(W, H, nvc.PixelFormat.NV12, nvc.PixelFormat.BGR, 0)
img = np.ndarray((H,W,3), dtype=np.uint8) nvDwn = nvc.PySurfaceDownloader(W, H, nvc.PixelFormat.BGR, 0)
img = np.ndarray((H,W,3), dtype=np.uint8)
if yuv: if yuv:
conv_yuv = nvc.PySurfaceConverter(W, H, nvc.PixelFormat.NV12, nvc.PixelFormat.YUV420, 0) conv_yuv = nvc.PySurfaceConverter(W, H, nvc.PixelFormat.NV12, nvc.PixelFormat.YUV420, 0)
@ -57,22 +60,24 @@ def decoder_nvidia(fn, vipc_server, vst, yuv=False):
rawSurface = decoder.DecodeSingleSurface() rawSurface = decoder.DecodeSingleSurface()
if rawSurface.Empty(): if rawSurface.Empty():
continue continue
convSurface = conv.Execute(rawSurface, cc1) if rgb:
nvDwn.DownloadSingleSurface(convSurface, img) convSurface = conv.Execute(rawSurface, cc1)
vipc_server.send(vst, img.flatten().data, cnt, 0, 0) nvDwn.DownloadSingleSurface(convSurface, img)
vipc_server.send(vst, img.flatten().data, cnt, 0, 0)
if yuv: if yuv:
convSurface = conv_yuv.Execute(rawSurface, cc1) convSurface = conv_yuv.Execute(rawSurface, cc1)
nvDwn_yuv.DownloadSingleSurface(convSurface, img_yuv) nvDwn_yuv.DownloadSingleSurface(convSurface, img_yuv)
vipc_server.send(vst+3, img_yuv.flatten().data, cnt, 0, 0) vipc_server.send(vst+3, img_yuv.flatten().data, cnt, 0, 0)
cnt += 1 cnt += 1
def decoder_ffmpeg(fn, vipc_server, vst, yuv=False): def decoder_ffmpeg(fn, vipc_server, vst, yuv=True, rgb=False):
import av # pylint: disable=import-error import av # pylint: disable=import-error
container = av.open(fn, options={"probesize": "32"}) container = av.open(fn, options={"probesize": "32"})
cnt = 0 cnt = 0
for frame in container.decode(video=0): for frame in container.decode(video=0):
img = frame.to_ndarray(format=av.video.format.VideoFormat('bgr24')) if rgb:
vipc_server.send(vst, img.flatten().data, cnt, 0, 0) img = frame.to_ndarray(format=av.video.format.VideoFormat('bgr24'))
vipc_server.send(vst, img.flatten().data, cnt, 0, 0)
if yuv: if yuv:
img_yuv = frame.to_ndarray(format=av.video.format.VideoFormat('yuv420p')) img_yuv = frame.to_ndarray(format=av.video.format.VideoFormat('yuv420p'))
vipc_server.send(vst+3, img_yuv.flatten().data, cnt, 0, 0) vipc_server.send(vst+3, img_yuv.flatten().data, cnt, 0, 0)
@ -83,7 +88,7 @@ if __name__ == "__main__":
parser = argparse.ArgumentParser(description='Decode video streams and broacast on VisionIPC') parser = argparse.ArgumentParser(description='Decode video streams and broacast on VisionIPC')
parser.add_argument("addr", help="Address of comma 3") parser.add_argument("addr", help="Address of comma 3")
parser.add_argument('--nvidia', action='store_true', help='Use nvidia instead of ffmpeg') parser.add_argument('--nvidia', action='store_true', help='Use nvidia instead of ffmpeg')
parser.add_argument('--yuv', action='store_true', help='Also broadcast YUV') parser.add_argument('--rgb', action='store_true', help='Also broadcast RGB')
parser.add_argument("--cams", default="0,1,2", help="Cameras to decode") parser.add_argument("--cams", default="0,1,2", help="Cameras to decode")
args = parser.parse_args() args = parser.parse_args()
@ -96,9 +101,9 @@ if __name__ == "__main__":
vipc_server = VisionIpcServer("camerad") vipc_server = VisionIpcServer("camerad")
for vst in cams.values(): for vst in cams.values():
vipc_server.create_buffers(vst, 4, True, W, H) if args.rgb:
if args.yuv: vipc_server.create_buffers(vst, 4, True, W, H)
vipc_server.create_buffers(vst+3, 4, False, W, H) vipc_server.create_buffers(vst+3, 4, False, W, H)
vipc_server.start_listener() vipc_server.start_listener()
for k,v in cams.items(): for k,v in cams.items():
@ -108,6 +113,6 @@ if __name__ == "__main__":
os.mkfifo(FIFO_NAME) os.mkfifo(FIFO_NAME)
multiprocessing.Process(target=writer, args=(FIFO_NAME, sys.argv[1], k)).start() multiprocessing.Process(target=writer, args=(FIFO_NAME, sys.argv[1], k)).start()
if args.nvidia: if args.nvidia:
multiprocessing.Process(target=decoder_nvidia, args=(FIFO_NAME, vipc_server, v, args.yuv)).start() multiprocessing.Process(target=decoder_nvidia, args=(FIFO_NAME, vipc_server, v, True, args.rgb)).start()
else: else:
multiprocessing.Process(target=decoder_ffmpeg, args=(FIFO_NAME, vipc_server, v, args.yuv)).start() multiprocessing.Process(target=decoder_ffmpeg, args=(FIFO_NAME, vipc_server, v, True, args.rgb)).start()

Loading…
Cancel
Save