Allow 1d audio data and 2d video data

Author: tikuma-lsuhsc

src/ffmpegio_plugin_numpy/__init__.py

@@ -29,7 +29,7 @@ def video_info(obj: ArrayLike) -> Tuple[Tuple[int, int, int], str]:
     :return: shape (height,width,components) and data type str
     :rtype: Tuple[Tuple[int, int, int], str]
     """
-    return obj.shape[-3:], obj.dtype.str
+    return obj.shape[-3:] if obj.ndim != 2 else [*obj.shape, 1], obj.dtype.str
 
 
 @hookimpl
@@ -41,7 +41,7 @@ def audio_info(obj: ArrayLike) -> Tuple[int, str]:
     :return: number of channels and sample data type in data type str
     :rtype: Tuple[Tuple[int], str]
     """
-    return obj.shape[-1:], obj.dtype.str
+    return obj.shape[-1:] if obj.ndim > 1 else [1], obj.dtype.str
 
 
 @hookimpl
@@ -71,7 +71,9 @@ def audio_bytes(obj: ArrayLike) -> memoryview:
 
 
 @hookimpl
-def bytes_to_video(b: bytes, dtype: str, shape: Tuple[int, int, int], squeeze: bool) -> 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

tests/test_base.py

@@ -21,7 +21,7 @@ def test_hooks():
 
     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))
@@ -46,8 +46,10 @@ def test_audio():
         pprint(probe.full_details(wavfile))
         fs1, x1 = audio.read(wavfile)
 
+        audio.write(wavfile, fs, x1[:, 0], show_log=True, overwrite=True)
+
     assert fs == fs1
-    assert x.shape==x1.shape
+    assert x.shape == x1.shape
     assert np.array_equal(x, x1)
 
 
@@ -61,6 +63,9 @@ def test_video():
         video.write(avifile, fs, F, pix_fmt="bgr24", vcodec="rawvideo", show_log=True)
         fs1, F1 = video.read(avifile)
 
+        # test 2-D input
+        video.write(avifile, fs, np.array(F[0, :, :, 0]), show_log=True, overwrite=True)
+
     assert fs == fs1
     assert np.array_equal(F, F1)