nv12: encoderd avoids a full frame copy (#24519)

* rgb to nv12

* nv12 works (w memcpy)

* correct now

* no copy

* fix nv12 with fast debayer

* reverts of unused stuff

* ui use nv12

* comment out thumbnails for now

* rebase fix

* dm read nv12

* model read nv12

* fix ffmpeg encoder

* thumbnails from nv12

* replay to nv12

* python framereader support nv12

* remove hardcoded frame/buffer sizes

* fix build

* ffmpeg encoder fix buffers

* small cleanup

* reduce power usage test

* fix cpu usage test

* fix snapshot

* fix loggerd test

* bump cereal

Co-authored-by: Comma Device <device@comma.ai>
Co-authored-by: Joost Wooning <jwooning@gmail.com>

tools/lib/framereader.py

@@ -185,20 +185,28 @@ def read_file_check_size(f, sz, cookie):
   return buff
 
 
-def rgb24toyuv420(rgb):
+def rgb24toyuv(rgb):
   yuv_from_rgb = np.array([[ 0.299     ,  0.587     ,  0.114      ],
                            [-0.14714119, -0.28886916,  0.43601035 ],
                            [ 0.61497538, -0.51496512, -0.10001026 ]])
   img = np.dot(rgb.reshape(-1, 3), yuv_from_rgb.T).reshape(rgb.shape)
 
-  y_len = img.shape[0] * img.shape[1]
-  uv_len = y_len // 4
+
 
   ys = img[:, :, 0]
   us = (img[::2, ::2, 1] + img[1::2, ::2, 1] + img[::2, 1::2, 1] + img[1::2, 1::2, 1]) / 4 + 128
   vs = (img[::2, ::2, 2] + img[1::2, ::2, 2] + img[::2, 1::2, 2] + img[1::2, 1::2, 2]) / 4 + 128
 
-  yuv420 = np.empty(y_len + 2 * uv_len, dtype=img.dtype)
+  return ys, us, vs
+
+
+def rgb24toyuv420(rgb):
+  ys, us, vs = rgb24toyuv(rgb)
+
+  y_len = rgb.shape[0] * rgb.shape[1]
+  uv_len = y_len // 4
+
+  yuv420 = np.empty(y_len + 2 * uv_len, dtype=rgb.dtype)
   yuv420[:y_len] = ys.reshape(-1)
   yuv420[y_len:y_len + uv_len] = us.reshape(-1)
   yuv420[y_len + uv_len:y_len + 2 * uv_len] = vs.reshape(-1)
@@ -206,6 +214,20 @@ def rgb24toyuv420(rgb):
   return yuv420.clip(0, 255).astype('uint8')
 
 
+def rgb24tonv12(rgb):
+  ys, us, vs = rgb24toyuv(rgb)
+
+  y_len = rgb.shape[0] * rgb.shape[1]
+  uv_len = y_len // 4
+
+  nv12 = np.empty(y_len + 2 * uv_len, dtype=rgb.dtype)
+  nv12[:y_len] = ys.reshape(-1)
+  nv12[y_len::2] = us.reshape(-1)
+  nv12[y_len+1::2] = vs.reshape(-1)
+
+  return nv12.clip(0, 255).astype('uint8')
+
+
 def decompress_video_data(rawdat, vid_fmt, w, h, pix_fmt):
   # using a tempfile is much faster than proc.communicate for some reason
 
@@ -237,6 +259,8 @@ def decompress_video_data(rawdat, vid_fmt, w, h, pix_fmt):
 
   if pix_fmt == "rgb24":
     ret = np.frombuffer(dat, dtype=np.uint8).reshape(-1, h, w, 3)
+  elif pix_fmt == "nv12":
+    ret = np.frombuffer(dat, dtype=np.uint8).reshape(-1, (h*w*3//2))
   elif pix_fmt == "yuv420p":
     ret = np.frombuffer(dat, dtype=np.uint8).reshape(-1, (h*w*3//2))
   elif pix_fmt == "yuv444p":
@@ -304,7 +328,7 @@ class RawFrameReader(BaseFrameReader):
     assert self.frame_count is not None
     assert num+count <= self.frame_count
 
-    if pix_fmt not in ("yuv420p", "rgb24"):
+    if pix_fmt not in ("nv12", "yuv420p", "rgb24"):
       raise ValueError(f"Unsupported pixel format {pix_fmt!r}")
 
     app = []
@@ -313,6 +337,8 @@ class RawFrameReader(BaseFrameReader):
       rgb_dat = self.load_and_debayer(dat)
       if pix_fmt == "rgb24":
         app.append(rgb_dat)
+      elif pix_fmt == "nv12":
+        app.append(rgb24tonv12(rgb_dat))
       elif pix_fmt == "yuv420p":
         app.append(rgb24toyuv420(rgb_dat))
       else:
@@ -329,7 +355,7 @@ class VideoStreamDecompressor:
     self.h = h
     self.pix_fmt = pix_fmt
 
-    if pix_fmt == "yuv420p":
+    if pix_fmt in ("nv12", "yuv420p"):
       self.out_size = w*h*3//2  # yuv420p
     elif pix_fmt in ("rgb24", "yuv444p"):
       self.out_size = w*h*3
@@ -387,6 +413,8 @@ class VideoStreamDecompressor:
           ret = np.frombuffer(dat, dtype=np.uint8).reshape((self.h, self.w, 3))
         elif self.pix_fmt == "yuv420p":
           ret = np.frombuffer(dat, dtype=np.uint8)
+        elif self.pix_fmt == "nv12":
+          ret = np.frombuffer(dat, dtype=np.uint8)
         elif self.pix_fmt == "yuv444p":
           ret = np.frombuffer(dat, dtype=np.uint8).reshape((3, self.h, self.w))
         else:
@@ -549,7 +577,7 @@ class GOPFrameReader(BaseFrameReader):
     if num + count > self.frame_count:
       raise ValueError(f"{num + count} > {self.frame_count}")
 
-    if pix_fmt not in ("yuv420p", "rgb24", "yuv444p"):
+    if pix_fmt not in ("nv12", "yuv420p", "rgb24", "yuv444p"):
       raise ValueError(f"Unsupported pixel format {pix_fmt!r}")
 
     ret = [self._get_one(num + i, pix_fmt) for i in range(count)]