Added squeeze arg to bytes_to_xxx

Author: tikuma-lsuhsc

ffmpegio

@@ -69,7 +69,7 @@ def audio_bytes(obj: ArrayLike) -> memoryview:
 
 
 @hookimpl
-def bytes_to_video(b: bytes, dtype: str, shape: Tuple[int, int, int]) -> ArrayLike:
+def bytes_to_video(b: bytes, dtype: str, shape: Tuple[int, int, int], squeeze: bool) -> ArrayLike:
     """convert bytes to rawvideo NumPy array
 
     :param b: byte data of arbitrary number of video frames
@@ -78,15 +78,18 @@ def bytes_to_video(b: bytes, dtype: str, shape: Tuple[int, int, int]) -> ArrayLi
     :type dtype: str
     :param size: frame dimension in pixels and number of color components (height, width, components)
     :type size: Tuple[int, int, int]
+    :param squeeze: True to remove all the singular dimensions
+    :type squeeze: bool
     :return: rawvideo frames
     :rtype: ArrayLike
     """
 
-    return np.frombuffer(b, dtype).reshape(-1, *shape)
+    x = np.frombuffer(b, dtype).reshape(-1, *shape)
+    return x.squeeze() if squeeze else x
 
 
 @hookimpl
-def bytes_to_audio(b: bytes, dtype: str, shape: Tuple[int]) -> ArrayLike:
+def bytes_to_audio(b: bytes, dtype: str, shape: Tuple[int], squeeze: bool) -> ArrayLike:
     """convert bytes to rawaudio NumPy array
 
     :param b: byte data of arbitrary number of video frames
@@ -95,8 +98,11 @@ def bytes_to_audio(b: bytes, dtype: str, shape: Tuple[int]) -> ArrayLike:
     :type dtype: str
     :param shape: number of audio channels
     :type shape: Tuple[int]
+    :param squeeze: True to remove all the singular dimensions
+    :type squeeze: bool
     :return: raw audio samples
     :rtype: ArrayLike
     """
 
-    return np.frombuffer(b, dtype).reshape(-1, *shape)
+    x = np.frombuffer(b, dtype).reshape(-1, *shape)
+    return x.squeeze() if squeeze else x

src/ffmpegio_plugin_numpy/__init__.py

@@ -11,17 +11,20 @@ def test_hooks():
     dtype = "|u1"
     shape = (2, 2, 3)
     b = b"\0" * utils.prod(shape)
-    data = hook.bytes_to_video(b=b, dtype=dtype, shape=shape)
+    data = hook.bytes_to_video(b=b, dtype=dtype, shape=shape, squeeze=False)
     assert memoryview(data).cast("b") == b
     assert data.dtype.str == dtype
     assert data.shape[1:] == shape
     assert hook.video_info(obj=data) == (shape, dtype)
     assert hook.video_bytes(obj=data).cast("b") == b
 
+    data = hook.bytes_to_video(b=b, dtype=dtype, shape=shape, squeeze=True)
+    assert data.shape == shape
+    
     dtype = "<f4"
     shape = (2,)
     b = b"\0" * (1024 * utils.prod(shape))
-    data = hook.bytes_to_audio(b=b, dtype=dtype, shape=shape)
+    data = hook.bytes_to_audio(b=b, dtype=dtype, shape=shape, squeeze=False)
     assert memoryview(data).cast("b") == b
     assert data.dtype.str == dtype
     assert data.shape[1:] == shape