From 534bb0187515878d5495e14d778bd63a330c4225 Mon Sep 17 00:00:00 2001
From: Mark Harfouche <mark.harfouche@gmail.com>
Date: Sat, 4 May 2024 10:13:10 -0400
Subject: [PATCH 1/2] Create fast path for returning YUVPlanes to numpy arrays

---
 av/video/frame.pyx | 22 ++++++++++++++++------
 av/video/plane.pxd |  4 ++++
 av/video/plane.pyx | 17 +++++++++++++++++
 3 files changed, 37 insertions(+), 6 deletions(-)

diff --git a/av/video/frame.pyx b/av/video/frame.pyx
index 6ff982491..13f1103b5 100644
--- a/av/video/frame.pyx
+++ b/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
 
@@ -295,11 +295,21 @@ cdef class VideoFrame(Frame):
         if frame.format.name in ("yuv420p", "yuvj420p"):
             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)
+            # Fast path for the case that the entire YUV data is contiguous
+            if (
+                frame.planes[0].line_size == frame.planes[0].width and
+                frame.planes[1].line_size == frame.planes[1].width and
+                frame.planes[2].line_size == frame.planes[2].width
+            ):
+                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]),
+                    useful_array(frame.planes[1]),
+                    useful_array(frame.planes[2])
+                )).reshape(-1, frame.width)
         elif frame.format.name in ("yuv444p", "yuvj444p"):
             return np.hstack((
                 useful_array(frame.planes[0]),
diff --git a/av/video/plane.pxd b/av/video/plane.pxd
index f9abf22b6..a74eea206 100644
--- a/av/video/plane.pxd
+++ b/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
diff --git a/av/video/plane.pyx b/av/video/plane.pyx
index 908b48716..08fddeb6c 100644
--- a/av/video/plane.pyx
+++ b/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]

From 08b3a7e619bb2fc92b343cd25761036602077b8d Mon Sep 17 00:00:00 2001
From: Mark Harfouche <mark.harfouche@gmail.com>
Date: Tue, 7 May 2024 07:09:19 -0400
Subject: [PATCH 2/2] Add a check for buffer continuity

---
 av/video/frame.pyx | 16 ++++++++++------
 1 file changed, 10 insertions(+), 6 deletions(-)

diff --git a/av/video/frame.pyx b/av/video/frame.pyx
index 13f1103b5..d5015421e 100644
--- a/av/video/frame.pyx
+++ b/av/video/frame.pyx
@@ -293,22 +293,26 @@ 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
+            y_plane, u_plane, v_plane = frame.planes[:3]
             # Fast path for the case that the entire YUV data is contiguous
             if (
-                frame.planes[0].line_size == frame.planes[0].width and
-                frame.planes[1].line_size == frame.planes[1].width and
-                frame.planes[2].line_size == frame.planes[2].width
+                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]),
-                    useful_array(frame.planes[1]),
-                    useful_array(frame.planes[2])
+                    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((