use nv12 for visionipc (#308)

* use nv12 for visionipc

* add uv_offet to python

* add more visionipc python functions

Co-authored-by: Comma Device <device@comma.ai>

visionipc/visionbuf.cc

@@ -14,14 +14,15 @@ void VisionBuf::init_rgb(size_t init_width, size_t init_height, size_t init_stri
   this->stride = init_stride;
 }
 
-void VisionBuf::init_yuv(size_t init_width, size_t init_height){
+void VisionBuf::init_yuv(size_t init_width, size_t init_height, size_t init_stride, size_t init_uv_offset){
   this->rgb = false;
   this->width = init_width;
   this->height = init_height;
+  this->stride = init_stride;
+  this->uv_offset = init_uv_offset;
 
   this->y = (uint8_t *)this->addr;
-  this->u = this->y + (this->width * this->height);
-  this->v = this->u + (this->width / 2 * this->height / 2);
+  this->uv = this->y + this->uv_offset;
 }
 
 

visionipc/visionbuf.h

@@ -36,11 +36,11 @@ class VisionBuf {
   size_t width = 0;
   size_t height = 0;
   size_t stride = 0;
+  size_t uv_offset = 0;
 
   // YUV
   uint8_t * y = nullptr;
-  uint8_t * u = nullptr;
-  uint8_t * v = nullptr;
+  uint8_t * uv = nullptr;
 
   // Visionipc
   uint64_t server_id = 0;
@@ -58,7 +58,7 @@ class VisionBuf {
   void import();
   void init_cl(cl_device_id device_id, cl_context ctx);
   void init_rgb(size_t width, size_t height, size_t stride);
-  void init_yuv(size_t width, size_t height);
+  void init_yuv(size_t width, size_t height, size_t stride, size_t uv_offset);
   int sync(int dir);
   int free();
 

visionipc/visionipc.pxd

@@ -16,6 +16,7 @@ cdef extern from "visionbuf.h":
     size_t width
     size_t height
     size_t stride
+    size_t uv_offset
     void set_frame_id(uint64_t id)
 
 cdef extern from "visionipc.h":
@@ -28,6 +29,7 @@ cdef extern from "visionipc_server.h":
   cdef cppclass VisionIpcServer:
     VisionIpcServer(string, void*, void*)
     void create_buffers(VisionStreamType, size_t, bool, size_t, size_t)
+    void create_buffers_with_sizes(VisionStreamType, size_t, bool, size_t, size_t, size_t, size_t, size_t)
     VisionBuf * get_buffer(VisionStreamType)
     void send(VisionBuf *, VisionIpcBufExtra *, bool)
     void start_listener()

visionipc/visionipc_client.cc

@@ -66,7 +66,7 @@ bool VisionIpcClient::connect(bool blocking){
     if (buffers[i].rgb) {
       buffers[i].init_rgb(buffers[i].width, buffers[i].height, buffers[i].stride);
     } else {
-      buffers[i].init_yuv(buffers[i].width, buffers[i].height);
+      buffers[i].init_yuv(buffers[i].width, buffers[i].height, buffers[i].stride, buffers[i].uv_offset);
     }
 
     if (device_id) buffers[i].init_cl(device_id, ctx);

visionipc/visionipc_pyx.pyx

@@ -31,9 +31,12 @@ cdef class VisionIpcServer:
   def __init__(self, string name):
     self.server = new cppVisionIpcServer(name, NULL, NULL)
 
-  def create_buffers(self,  VisionStreamType tp, size_t num_buffers, bool rgb, size_t width, size_t height):
+  def create_buffers(self, VisionStreamType tp, size_t num_buffers, bool rgb, size_t width, size_t height):
     self.server.create_buffers(tp, num_buffers, rgb, width, height)
 
+  def create_buffers_with_sizes(self, VisionStreamType tp, size_t num_buffers, bool rgb, size_t width, size_t height, size_t size, size_t stride, size_t uv_offset):
+    self.server.create_buffers_with_sizes(tp, num_buffers, rgb, width, height, size, stride, uv_offset)
+
   def send(self, VisionStreamType tp, const unsigned char[:] data, uint32_t frame_id=0, uint64_t timestamp_sof=0, uint64_t timestamp_eof=0):
     cdef cppVisionBuf * buf = self.server.get_buffer(tp)
 
@@ -79,6 +82,10 @@ cdef class VisionIpcClient:
   def stride(self):
     return None if not self.buf else self.buf.stride
 
+  @property
+  def uv_offset(self):
+    return None if not self.buf else self.buf.uv_offset
+
   def recv(self, int timeout_ms=100):
     self.buf = self.client.recv(NULL, timeout_ms)
     if not self.buf:

visionipc/visionipc_server.cc

@@ -35,7 +35,8 @@ void VisionIpcServer::create_buffers(VisionStreamType type, size_t num_buffers,
   int aligned_w = 0, aligned_h = 0;
 
   size_t size = 0;
-  size_t stride = 0; // Only used for RGB
+  size_t stride = 0;
+  size_t uv_offset = 0;
 
   if (rgb) {
     visionbuf_compute_aligned_width_and_height(width, height, &aligned_w, &aligned_h);
@@ -43,8 +44,14 @@ void VisionIpcServer::create_buffers(VisionStreamType type, size_t num_buffers,
     stride = aligned_w * 3;
   } else {
     size = width * height * 3 / 2;
+    stride = width;
+    uv_offset = width * height;
   }
 
+  create_buffers_with_sizes(type, num_buffers, rgb, width, height, size, stride, uv_offset);
+}
+
+void VisionIpcServer::create_buffers_with_sizes(VisionStreamType type, size_t num_buffers, bool rgb, size_t width, size_t height, size_t size, size_t stride, size_t uv_offset) {
   // Create map + alloc requested buffers
   for (size_t i = 0; i < num_buffers; i++){
     VisionBuf* buf = new VisionBuf();
@@ -54,7 +61,7 @@ void VisionIpcServer::create_buffers(VisionStreamType type, size_t num_buffers,
 
     if (device_id) buf->init_cl(device_id, ctx);
 
-    rgb ? buf->init_rgb(width, height, stride) : buf->init_yuv(width, height);
+    rgb ? buf->init_rgb(width, height, stride) : buf->init_yuv(width, height, stride, uv_offset);
 
     buffers[type].push_back(buf);
   }

visionipc/visionipc_server.h

@@ -37,6 +37,7 @@ class VisionIpcServer {
   VisionBuf * get_buffer(VisionStreamType type);
 
   void create_buffers(VisionStreamType type, size_t num_buffers, bool rgb, size_t width, size_t height);
+  void create_buffers_with_sizes(VisionStreamType type, size_t num_buffers, bool rgb, size_t width, size_t height, size_t size, size_t stride, size_t uv_offset);
   void send(VisionBuf * buf, VisionIpcBufExtra * extra, bool sync=true);
   void start_listener();
 };