From 26c45da33c0f47e7fe03c1666bdd91f011dbcf71 Mon Sep 17 00:00:00 2001 From: Mark Harfouche Date: Sat, 4 May 2024 10:13:10 -0400 Subject: [PATCH] Create fast path for returning YUVPlanes to numpy arrays --- av/video/frame.pyx | 22 ++++++++++++++++------ 1 file changed, 16 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]),