From 10d0f5fbeb17cf5391e174315f495f532e3c5282 Mon Sep 17 00:00:00 2001
From: Jonathan Drolet <jonathan.drolet@riedel.net>
Date: Tue, 13 Sep 2022 14:38:29 -0400
Subject: [PATCH] Add support in VideoFrame ndarray for gbrpf32 format
---
av/video/frame.pyx | 15 +++++++++++++++
tests/test_videoframe.py | 18 ++++++++++++++++++
2 files changed, 33 insertions(+)
diff --git a/av/video/frame.pyx b/av/video/frame.pyx
index 67cf553..ac226fa 100644
--- a/av/video/frame.pyx
+++ b/av/video/frame.pyx
@@ -285,6 +285,12 @@ cdef class VideoFrame(Frame):
array[:, :, 1] = useful_array(frame.planes[0], 2, "uint16").reshape(-1, frame.width)
array[:, :, 2] = useful_array(frame.planes[1], 2, "uint16").reshape(-1, frame.width)
return byteswap_array(array, frame.format.name.endswith('be'))
+ elif frame.format.name in ('gbrpf32be', 'gbrpf32le'):
+ array = np.empty((frame.height, frame.width, 3), dtype="float32")
+ array[:, :, 0] = useful_array(frame.planes[2], 4, "float32").reshape(-1, frame.width)
+ array[:, :, 1] = useful_array(frame.planes[0], 4, "float32").reshape(-1, frame.width)
+ array[:, :, 2] = useful_array(frame.planes[1], 4, "float32").reshape(-1, frame.width)
+ return byteswap_array(array, frame.format.name.endswith('be'))
elif frame.format.name in ('rgb24', 'bgr24'):
return useful_array(frame.planes[0], 3).reshape(frame.height, frame.width, -1)
elif frame.format.name in ('argb', 'rgba', 'abgr', 'bgra'):
@@ -389,6 +395,15 @@ cdef class VideoFrame(Frame):
copy_array_to_plane(byteswap_array(array[:, :, 2], format.endswith('be')), frame.planes[1], 2)
copy_array_to_plane(byteswap_array(array[:, :, 0], format.endswith('be')), frame.planes[2], 2)
return frame
+ elif format in ('gbrpf32be', 'gbrpf32le'):
+ check_ndarray(array, 'float32', 3)
+ check_ndarray_shape(array, array.shape[2] == 3)
+
+ frame = VideoFrame(array.shape[1], array.shape[0], format)
+ copy_array_to_plane(byteswap_array(array[:, :, 1], format.endswith('be')), frame.planes[0], 4)
+ copy_array_to_plane(byteswap_array(array[:, :, 2], format.endswith('be')), frame.planes[1], 4)
+ copy_array_to_plane(byteswap_array(array[:, :, 0], format.endswith('be')), frame.planes[2], 4)
+ return frame
elif format in ('rgb24', 'bgr24'):
check_ndarray(array, 'uint8', 3)
check_ndarray_shape(array, array.shape[2] == 3)
diff --git a/tests/test_videoframe.py b/tests/test_videoframe.py
index a322c1e..3d354f0 100644
--- a/tests/test_videoframe.py
+++ b/tests/test_videoframe.py
@@ -295,6 +295,24 @@ class TestVideoFrameNdarray(TestCase):
self.assertEqual(frame.format.name, format)
self.assertNdarraysEqual(frame.to_ndarray(), array)
+ def test_ndarray_gbrpf32(self):
+ array = numpy.random.random_sample(size=(480, 640, 3)).astype(numpy.float32)
+ for format in ["gbrpf32be", "gbrpf32le"]:
+ frame = VideoFrame.from_ndarray(array, format=format)
+ self.assertEqual(frame.width, 640)
+ self.assertEqual(frame.height, 480)
+ self.assertEqual(frame.format.name, format)
+ self.assertNdarraysEqual(frame.to_ndarray(), array)
+
+ def test_ndarray_gbrpf32_align(self):
+ array = numpy.random.random_sample(size=(238, 318, 3)).astype(numpy.float32)
+ for format in ["gbrpf32be", "gbrpf32le"]:
+ frame = VideoFrame.from_ndarray(array, format=format)
+ self.assertEqual(frame.width, 318)
+ self.assertEqual(frame.height, 238)
+ self.assertEqual(frame.format.name, format)
+ self.assertNdarraysEqual(frame.to_ndarray(), array)
+
def test_ndarray_yuv420p(self):
array = numpy.random.randint(0, 256, size=(720, 640), dtype=numpy.uint8)
frame = VideoFrame.from_ndarray(array, format="yuv420p")
--
GitLab