Create fast path for returning YUVPlanes to numpy arrays

av/video/frame.pyx

@@ -6,7 +6,7 @@ from av.enum cimport define_enum
 from av.error cimport err_check
 from av.utils cimport check_ndarray, check_ndarray_shape
 from av.video.format cimport get_pix_fmt, get_video_format
-from av.video.plane cimport VideoPlane
+from av.video.plane cimport VideoPlane, YUVPlanes
 
 import warnings
 
@@ -293,13 +293,27 @@ cdef class VideoFrame(Frame):
         import numpy as np
 
         if frame.format.name in ("yuv420p", "yuvj420p"):
+            bytes_per_pixel = 1
             assert frame.width % 2 == 0
             assert frame.height % 2 == 0
-            return np.hstack((
-                useful_array(frame.planes[0]),
-                useful_array(frame.planes[1]),
-                useful_array(frame.planes[2])
-            )).reshape(-1, frame.width)
+            y_plane, u_plane, v_plane = frame.planes[:3]
+            # Fast path for the case that the entire YUV data is contiguous
+            if (
+                y_plane.line_size == y_plane.width and
+                u_plane.line_size == u_plane.width and
+                v_plane.line_size == v_plane.width and
+                y_plane.buffer_ptr + y_plane.buffer_size * bytes_per_pixel == u_plane.buffer_ptr and
+                u_plane.buffer_ptr + u_plane.buffer_size * bytes_per_pixel == v_plane.buffer_ptr
+            ):
+                yuv_planes = YUVPlanes(frame, 0)
+                return useful_array(yuv_planes).reshape(frame.height * 3 // 2, frame.width)
+            else:
+                # Otherwise, we need to copy the data through the use of np.hstack
+                return np.hstack((
+                    useful_array(frame.planes[0], bytes_per_pixel=bytes_per_pixel),
+                    useful_array(frame.planes[1], bytes_per_pixel=bytes_per_pixel),
+                    useful_array(frame.planes[2], bytes_per_pixel=bytes_per_pixel)
+                )).reshape(-1, frame.width)
         elif frame.format.name in ("yuv444p", "yuvj444p"):
             return np.hstack((
                 useful_array(frame.planes[0]),

av/video/plane.pxd

@@ -6,3 +6,7 @@ cdef class VideoPlane(Plane):
 
     cdef readonly size_t buffer_size
     cdef readonly unsigned int width, height
+
+
+cdef class YUVPlanes(VideoPlane):
+    pass

av/video/plane.pyx

@@ -35,3 +35,20 @@ cdef class VideoPlane(Plane):
         :type: int
         """
         return self.frame.ptr.linesize[self.index]
+
+
+cdef class YUVPlanes(VideoPlane):
+    def __cinit__(self, VideoFrame frame, int index):
+        if index != 0:
+            raise RuntimeError("YUVPlanes only supports index 0")
+        if frame.format.name not in ['yuvj420p', 'yuv420p']:
+            raise RuntimeError("YUVPlane only supports yuv420p and yuvj420p")
+        if frame.ptr.linesize[0] < 0:
+            raise RuntimeError("YUVPlane only supports positive linesize")
+        self.width = frame.width
+        self.height = frame.height * 3 // 2
+        self.buffer_size = self.height *  abs(self.frame.ptr.linesize[0])
+        self.frame = frame
+
+    cdef void* _buffer_ptr(self):
+        return self.frame.ptr.extended_data[self.index]