Skip to content

Commit

Permalink
Create a dedicate code path to convert YUV frames to numpy for speed
Browse files Browse the repository at this point in the history
This avoids the use of hstack, which inevitiably copies the data to a
new memory locaiton.

The speed up is small, but measurable.

I can go from 185 fps decoding of a ~3000 x 2000  video to 200 fps
decoding.
  • Loading branch information
hmaarrfk committed Jul 23, 2023
1 parent 3fbb392 commit a78a054
Show file tree
Hide file tree
Showing 3 changed files with 37 additions and 6 deletions.
22 changes: 16 additions & 6 deletions av/video/frame.pyx
Original file line number Diff line number Diff line change
Expand Up @@ -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 VideoFormat, get_pix_fmt, get_video_format
from av.video.plane cimport VideoPlane
from av.video.plane cimport VideoPlane, YUVPlanes


cdef object _cinit_bypass_sentinel
Expand Down Expand Up @@ -265,11 +265,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 == 'yuyv422':
assert frame.width % 2 == 0
assert frame.height % 2 == 0
Expand Down
4 changes: 4 additions & 0 deletions av/video/plane.pxd
Original file line number Diff line number Diff line change
Expand Up @@ -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
17 changes: 17 additions & 0 deletions av/video/plane.pyx
Original file line number Diff line number Diff line change
Expand Up @@ -37,3 +37,20 @@ cdef class VideoPlane(Plane):
"""
def __get__(self):
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]

0 comments on commit a78a054

Please sign in to comment.