diff --git a/av/video/frame.pyx b/av/video/frame.pyx index 42ec9dea2..67cf55382 100644 --- a/av/video/frame.pyx +++ b/av/video/frame.pyx @@ -310,6 +310,11 @@ cdef class VideoFrame(Frame): image = useful_array(frame.planes[0]).reshape(frame.height, frame.width) palette = np.frombuffer(frame.planes[1], 'i4').astype('>i4').reshape(-1, 1).view(np.uint8) return image, palette + elif frame.format.name == 'nv12': + return np.hstack(( + useful_array(frame.planes[0]), + useful_array(frame.planes[1], 2) + )).reshape(-1, frame.width) else: raise ValueError('Conversion to numpy array with format `%s` is not yet supported' % frame.format.name) @@ -409,6 +414,17 @@ cdef class VideoFrame(Frame): frame = VideoFrame(array.shape[1], array.shape[0], format) copy_array_to_plane(byteswap_array(array, format == 'rgba64be'), frame.planes[0], 8) return frame + elif format == 'nv12': + check_ndarray(array, 'uint8', 2) + check_ndarray_shape(array, array.shape[0] % 3 == 0) + check_ndarray_shape(array, array.shape[1] % 2 == 0) + + frame = VideoFrame(array.shape[1], (array.shape[0] * 2) // 3, format) + uv_start = frame.width * frame.height + flat = array.reshape(-1) + copy_array_to_plane(flat[:uv_start], frame.planes[0], 1) + copy_array_to_plane(flat[uv_start:], frame.planes[1], 2) + return frame else: raise ValueError('Conversion from numpy array with format `%s` is not yet supported' % format) diff --git a/tests/test_videoframe.py b/tests/test_videoframe.py index f037326db..a322c1e99 100644 --- a/tests/test_videoframe.py +++ b/tests/test_videoframe.py @@ -440,6 +440,22 @@ def test_ndarray_pal8(self): self.assertNdarraysEqual(returned[0], array) self.assertNdarraysEqual(returned[1], palette) + def test_ndarray_nv12(self): + array = numpy.random.randint(0, 256, size=(720, 640), dtype=numpy.uint8) + frame = VideoFrame.from_ndarray(array, format="nv12") + self.assertEqual(frame.width, 640) + self.assertEqual(frame.height, 480) + self.assertEqual(frame.format.name, "nv12") + self.assertNdarraysEqual(frame.to_ndarray(), array) + + def test_ndarray_nv12_align(self): + array = numpy.random.randint(0, 256, size=(357, 318), dtype=numpy.uint8) + frame = VideoFrame.from_ndarray(array, format="nv12") + self.assertEqual(frame.width, 318) + self.assertEqual(frame.height, 238) + self.assertEqual(frame.format.name, "nv12") + self.assertNdarraysEqual(frame.to_ndarray(), array) + class TestVideoFrameTiming(TestCase): def test_reformat_pts(self):