diff --git a/imageio_ffmpeg/_io.py b/imageio_ffmpeg/_io.py index 0cb3b00..5e43eef 100644 --- a/imageio_ffmpeg/_io.py +++ b/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 diff --git a/tests/test_io.py b/tests/test_io.py index d496919..ae5cf96 100644 --- a/tests/test_io.py +++ b/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