diff --git a/av/video/codeccontext.pyx b/av/video/codeccontext.pyx index 33efc5d54..51103b787 100644 --- a/av/video/codeccontext.pyx +++ b/av/video/codeccontext.pyx @@ -161,3 +161,15 @@ cdef class VideoCodecContext(CodecContext): property coded_height: def __get__(self): return self.ptr.coded_height + + property max_b_frames: + """ + The maximum run of consecutive B frames when encoding a video. + + :type: int + """ + def __get__(self): + return self.ptr.max_b_frames + + def __set__(self, value): + self.ptr.max_b_frames = value diff --git a/tests/test_encode.py b/tests/test_encode.py index a293cbf83..e79bf539f 100644 --- a/tests/test_encode.py +++ b/tests/test_encode.py @@ -1,7 +1,10 @@ from fractions import Fraction from unittest import SkipTest +import io import math +import numpy as np + from av import AudioFrame, VideoFrame from av.audio.stream import AudioStream from av.video.stream import VideoStream @@ -291,3 +294,84 @@ def test_set_id_and_time_base(self): self.assertEqual(stream.time_base, None) stream.time_base = Fraction(1, 48000) self.assertEqual(stream.time_base, Fraction(1, 48000)) + + +def encode_file_with_max_b_frames(max_b_frames): + """ + Create an encoded video file (or file-like object) with the given + maximum run of B frames. + + max_b_frames: non-negative integer which is the maximum allowed run + of consecutive B frames. + + Returns: a file-like object. + """ + # Create a video file that is entirely arbitrary, but with the passed + # max_b_frames parameter. + file = io.BytesIO() + container = av.open(file, mode="w", format="mp4") + stream = container.add_stream("h264", rate=30) + stream.width = 640 + stream.height = 480 + stream.pix_fmt = "yuv420p" + stream.codec_context.gop_size = 15 + stream.codec_context.max_b_frames = max_b_frames + + for i in range(50): + array = np.empty((stream.height, stream.width, 3), dtype=np.uint8) + # This appears to hit a complexity "sweet spot" that makes the codec + # want to use B frames. + array[:, :] = (i, 0, 255 - i) + frame = av.VideoFrame.from_ndarray(array, format="rgb24") + for packet in stream.encode(frame): + container.mux(packet) + + for packet in stream.encode(): + container.mux(packet) + + container.close() + file.seek(0) + + return file + + +def max_b_frame_run_in_file(file): + """ + Count the maximum run of B frames in a file (or file-like object). + + file: the file or file-like object in which to count the maximum run + of B frames. The file should contain just one video stream. + + Returns: non-negative integer which is the maximum B frame run length. + """ + container = av.open(file) + stream = container.streams.video[0] + + max_b_frame_run = 0 + b_frame_run = 0 + for packet in container.demux(stream): + for frame in packet.decode(): + if frame.pict_type == av.video.frame.PictureType.B: + b_frame_run += 1 + else: + max_b_frame_run = max(max_b_frame_run, b_frame_run) + b_frame_run = 0 + + # Outside chance that the longest run was at the end of the file. + max_b_frame_run = max(max_b_frame_run, b_frame_run) + + container.close() + + return max_b_frame_run + + +class TestMaxBFrameEncoding(TestCase): + def test_max_b_frames(self): + """ + Test that we never get longer runs of B frames than we asked for with + the max_b_frames property. + """ + for max_b_frames in range(4): + file = encode_file_with_max_b_frames(max_b_frames) + actual_max_b_frames = max_b_frame_run_in_file(file) + self.assertTrue(actual_max_b_frames <= max_b_frames)