added bits_per_pixel parameter to read_frames

imageio_ffmpeg/_io.py

@@ -60,7 +60,7 @@ def count_frames_and_secs(path):
     raise RuntimeError("Could not get number of frames")  # pragma: no cover
 
 
-def read_frames(path, pix_fmt="rgb24", bpp=3, input_params=None, output_params=None):
+def read_frames(path, pix_fmt="rgb24", bpp=None, input_params=None, output_params=None, bits_per_pixel=None):
     """
     Create a generator to iterate over the frames in a video file.
     
@@ -99,6 +99,9 @@ def read_frames(path, pix_fmt="rgb24", bpp=3, input_params=None, output_params=N
             This depends on the given pix_fmt. Default is 3 (RGB).
         input_params (list): Additional ffmpeg input command line parameters.
         output_params (list): Additional ffmpeg output command line parameters.
+        bits_per_pixel (int): The number of bits per pixel in the output frames.
+            This depends on the given pix_fmt. Yuv format can have 12 bits per pixel,
+            so bpp attribute useless
     """
 
     # ----- Input args
@@ -108,11 +111,12 @@ def read_frames(path, pix_fmt="rgb24", bpp=3, input_params=None, output_params=N
 
     pix_fmt = pix_fmt or "rgb24"
     bpp = bpp or 3
+    bits_per_pixel = bits_per_pixel or bpp * 8
     input_params = input_params or []
     output_params = output_params or []
 
     assert isinstance(pix_fmt, str), "pix_fmt must be a string"
-    assert isinstance(bpp, int), "bpp must be an int"
+    assert isinstance(bits_per_pixel, int), "bpp and bits_per_pixel must be an int"
     assert isinstance(input_params, list), "input_params must be a list"
     assert isinstance(output_params, list), "output_params must be a list"
 
@@ -157,15 +161,18 @@ def read_frames(path, pix_fmt="rgb24", bpp=3, input_params=None, output_params=N
         # ----- Read frames
 
         w, h = meta["size"]
-        framesize = w * h * bpp
+        framesize_bits = w * h * bits_per_pixel
+        framesize_bytes = framesize_bits / 8
+        assert framesize_bytes.is_integer(), "incorrect bits_per_pixel, framesize in bytes must be an int"
+        framesize_bytes = int(framesize_bytes)
         framenr = 0
 
         while True:
             framenr += 1
             try:
                 bb = bytes()
-                while len(bb) < framesize:
-                    extra_bytes = p.stdout.read(framesize - len(bb))
+                while len(bb) < framesize_bytes:
+                    extra_bytes = p.stdout.read(framesize_bytes - len(bb))
                     if not extra_bytes:
                         if len(bb) == 0:
                             return

tests/test_io.py

@@ -16,15 +16,20 @@
 
 
 test_dir = tempfile.gettempdir()
-test_url = "https://raw.githubusercontent.com/imageio/imageio-binaries/master/images/cockatoo.mp4"
+test_url1 = "https://raw.githubusercontent.com/imageio/imageio-binaries/master/images/cockatoo.mp4"
+test_url2 = "https://raw.githubusercontent.com/imageio/imageio-binaries/master/images/realshort.mp4"
 test_file1 = os.path.join(test_dir, "cockatoo.mp4")
 test_file2 = os.path.join(test_dir, "test.mp4")
+test_file3 = os.path.join(test_dir, "realshort.mp4")
 
 
 def setup_module():
-    bb = urlopen(test_url, timeout=5).read()
+    bb = urlopen(test_url1, timeout=5).read()
     with open(test_file1, "wb") as f:
         f.write(bb)
+    bb = urlopen(test_url2, timeout=5).read()
+    with open(test_file3, "wb") as f:
+        f.write(bb)
 
 
 def test_ffmpeg_version():
@@ -108,6 +113,28 @@ def test_reading4():
     assert "ffmpeg version" in msg  # The log is included
 
 
+def test_reading5():
+    # Same as 1, but using other pixel format and bits_per_pixel
+    bits_per_pixel = 12
+    bits_per_bytes = 8
+    gen = imageio_ffmpeg.read_frames(test_file3, pix_fmt="yuv420p", bits_per_pixel=bits_per_pixel)
+
+    meta = gen.__next__()
+    assert isinstance(meta, dict)
+    for key in ("size", "fps", "duration"):
+        assert key in meta
+
+    # Read frames
+    framesize = meta["size"][0] * meta["size"][1] * bits_per_pixel / bits_per_bytes
+    assert framesize == 320 * 240 * bits_per_pixel / bits_per_bytes
+    count = 0
+    for frame in gen:
+        assert isinstance(frame, bytes) and len(frame) == framesize
+        count += 1
+
+    assert count == 36
+
+
 def test_reading_invalid_video():
     """
     Check whether invalid video is