Expose VisionBuf allocation to Python

msgq/visionipc/visionipc.pxd

@@ -4,7 +4,7 @@
 from libcpp.string cimport string
 from libcpp.vector cimport vector
 from libcpp.set cimport set
-from libc.stdint cimport uint32_t, uint64_t
+from libc.stdint cimport uint8_t, uint32_t, uint64_t
 from libcpp cimport bool, int
 
 cdef extern from "msgq/visionipc/visionbuf.h":
@@ -13,13 +13,21 @@ cdef extern from "msgq/visionipc/visionbuf.h":
 
   cdef cppclass VisionBuf:
     void * addr
+    uint8_t * y
+    uint8_t * uv
     int fd
     size_t len
+    size_t mmap_len
     size_t width
     size_t height
     size_t stride
     size_t uv_offset
     size_t idx
+    void allocate(size_t)
+    void init_yuv(size_t, size_t, size_t, size_t)
+    int sync(int)
+    int free()
+    uint64_t get_frame_id()
     void set_frame_id(uint64_t id)
 
 cdef extern from "msgq/visionipc/visionipc.h":

msgq/visionipc/visionipc_pyx.pxd

@@ -5,6 +5,7 @@ from .visionipc cimport VisionBuf as cppVisionBuf
 
 cdef class VisionBuf:
   cdef cppVisionBuf * buf
+  cdef bint owner
 
   @staticmethod
   cdef create(cppVisionBuf*)

msgq/visionipc/visionipc_pyx.pyx

@@ -6,7 +6,7 @@ import numpy as np
 cimport numpy as cnp
 from cython.view cimport array
 from libc.string cimport memcpy
-from libc.stdint cimport uint32_t, uint64_t
+from libc.stdint cimport uint8_t, uint32_t, uint64_t, uintptr_t
 from libcpp cimport bool
 from libcpp.string cimport string
 
@@ -29,16 +29,47 @@ cpdef enum VisionStreamType:
 
 
 cdef class VisionBuf:
+  def __cinit__(self):
+    self.buf = new cppVisionBuf()
+    self.owner = True
+
   @staticmethod
   cdef create(cppVisionBuf * cbuf):
     buf = VisionBuf()
+    del buf.buf
     buf.buf = cbuf
+    buf.owner = False
     return buf
 
+  def __dealloc__(self):
+    if self.owner and self.buf != NULL:
+      del self.buf
+
   @property
   def data(self):
     return np.asarray(<cnp.uint8_t[:self.buf.len]> self.buf.addr)
 
+  @property
+  def y(self):
+    return np.asarray(<cnp.uint8_t[:self.buf.stride * self.buf.height]> <uint8_t *>self.buf.y)
+
+  @property
+  def uv(self):
+    cdef size_t uv_len = self.buf.len - self.buf.uv_offset
+    return np.asarray(<cnp.uint8_t[:uv_len]> <uint8_t *>self.buf.uv)
+
+  @property
+  def addr(self):
+    return <uintptr_t>self.buf.addr
+
+  @property
+  def y_addr(self):
+    return <uintptr_t>self.buf.y
+
+  @property
+  def uv_addr(self):
+    return <uintptr_t>self.buf.uv
+
   @property
   def width(self):
     return self.buf.width
@@ -63,6 +94,24 @@ cdef class VisionBuf:
   def fd(self):
     return self.buf.fd
 
+  def allocate(self, size_t length):
+    self.buf.allocate(length)
+
+  def init_yuv(self, size_t width, size_t height, size_t stride, size_t uv_offset):
+    self.buf.init_yuv(width, height, stride, uv_offset)
+
+  def sync(self, int direction):
+    return self.buf.sync(direction)
+
+  def free(self):
+    return self.buf.free()
+
+  def get_frame_id(self):
+    return self.buf.get_frame_id()
+
+  def set_frame_id(self, uint64_t frame_id):
+    self.buf.set_frame_id(frame_id)
+
 
 cdef class VisionIpcServer:
   cdef cppVisionIpcServer * server