Skip to content

Commit

Permalink
Add max_b_frames property to codec context
Browse files Browse the repository at this point in the history
Allows an application to control the maximum number of B frames,
irrespective of what other settings or preset is used. A test is
also added, to check that it works as intended.
  • Loading branch information
davidplowman committed Nov 10, 2023
1 parent a4bde60 commit 1e3053a
Show file tree
Hide file tree
Showing 2 changed files with 96 additions and 0 deletions.
12 changes: 12 additions & 0 deletions av/video/codeccontext.pyx
Original file line number Diff line number Diff line change
Expand Up @@ -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
84 changes: 84 additions & 0 deletions tests/test_encode.py
Original file line number Diff line number Diff line change
@@ -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
Expand Down Expand Up @@ -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)

0 comments on commit 1e3053a

Please sign in to comment.