diff --git a/av/video/frame.pyx b/av/video/frame.pyx index ac226fa0f..972abb779 100644 --- a/av/video/frame.pyx +++ b/av/video/frame.pyx @@ -269,6 +269,12 @@ cdef class VideoFrame(Frame): 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]), + useful_array(frame.planes[1]), + useful_array(frame.planes[2]) + )).reshape(-1, frame.height, frame.width) elif frame.format.name == 'yuyv422': assert frame.width % 2 == 0 assert frame.height % 2 == 0 @@ -372,6 +378,16 @@ cdef class VideoFrame(Frame): copy_array_to_plane(flat[u_start:v_start], frame.planes[1], 1) copy_array_to_plane(flat[v_start:], frame.planes[2], 1) return frame + elif format in ('yuv444p', 'yuvj444p'): + check_ndarray(array, 'uint8', 3) + check_ndarray_shape(array, array.shape[0] == 3) + + frame = VideoFrame(array.shape[2], array.shape[1], format) + array = array.reshape(3, -1) + copy_array_to_plane(array[0], frame.planes[0], 1) + copy_array_to_plane(array[1], frame.planes[1], 1) + copy_array_to_plane(array[2], frame.planes[2], 1) + return frame elif format == 'yuyv422': check_ndarray(array, 'uint8', 3) check_ndarray_shape(array, array.shape[0] % 2 == 0) diff --git a/tests/test_videoframe.py b/tests/test_videoframe.py index 283343bb6..0d36616e1 100644 --- a/tests/test_videoframe.py +++ b/tests/test_videoframe.py @@ -345,6 +345,22 @@ def test_ndarray_yuyv422(self): self.assertEqual(frame.format.name, "yuyv422") self.assertNdarraysEqual(frame.to_ndarray(), array) + def test_ndarray_yuv444p(self): + array = numpy.random.randint(0, 256, size=(3, 480, 640), dtype=numpy.uint8) + frame = VideoFrame.from_ndarray(array, format="yuv444p") + self.assertEqual(frame.width, 640) + self.assertEqual(frame.height, 480) + self.assertEqual(frame.format.name, "yuv444p") + self.assertNdarraysEqual(frame.to_ndarray(), array) + + def test_ndarray_yuvj444p(self): + array = numpy.random.randint(0, 256, size=(3, 480, 640), dtype=numpy.uint8) + frame = VideoFrame.from_ndarray(array, format="yuvj444p") + self.assertEqual(frame.width, 640) + self.assertEqual(frame.height, 480) + self.assertEqual(frame.format.name, "yuvj444p") + self.assertNdarraysEqual(frame.to_ndarray(), array) + def test_ndarray_yuyv422_align(self): array = numpy.random.randint(0, 256, size=(238, 318, 2), dtype=numpy.uint8) frame = VideoFrame.from_ndarray(array, format="yuyv422")