BUG: Upgrade pyav plugin to support av v14 (#1112)

shaform · FirefoxMetzger · web-flow · commit 08d2cc145495 · 2025-01-18T00:32:53.000-08:00
Note: av v14 no longer supports the "libx264" codec. Use "h264" instead :)

---------

Co-authored-by: Sebastian Wallk&ouml;tter &lt;sebastian@wallkoetter.net&gt;
diff --git a/.github/workflows/pypy_tests.yml b/.github/workflows/pypy_tests.yml
@@ -5,19 +5,18 @@ on:
 
 jobs:
   pypy:
-    name: "${{ matrix.os }} / ${{ matrix.pyversion }}"
+    name: "${{ matrix.os }} / pypy-3.10"
     runs-on: ${{ matrix.os }}
     strategy:
       fail-fast: true
       matrix:
-        os: ["ubuntu-latest", "macos-12"]
-        pyversion: ["pypy-3.9", "pypy-3.10"]
+        os: ["ubuntu-latest", "macos-latest"]
     steps:
       - uses: actions/checkout@v4
       - name: Set up pypy
         uses: actions/setup-python@v5
         with:
-          python-version: "${{ matrix.pyversion }}"
+          python-version: "pypy-3.10"
       - name: MacOS Numpy Fix
         if: runner.os == 'macOS'
         run: |
diff --git a/imageio/plugins/pyav.py b/imageio/plugins/pyav.py
@@ -177,11 +177,12 @@
 
 from fractions import Fraction
 from math import ceil
-from typing import Any, Dict, List, Optional, Tuple, Union, Generator
+from typing import Any, Dict, Generator, List, Optional, Tuple, Union
 
 import av
 import av.filter
 import numpy as np
+from av.codec.context import Flags
 from numpy.lib.stride_tricks import as_strided
 
 from ..core import Request
@@ -258,6 +259,39 @@ def _get_frame_shape(frame: av.VideoFrame) -> Tuple[int, ...]:
     return tuple(shape)
 
 
+def _get_frame_type(picture_type: int) -> str:
+    """Return a human-readable name for provided picture type
+
+    Parameters
+    ----------
+    picture_type : int
+        The picture type extracted from Frame.pict_type
+
+    Returns
+    -------
+    picture_name : str
+        A human readable name of the picture type
+
+    """
+
+    if not isinstance(picture_type, int):
+        # old pyAV versions send an enum, not an int
+        return picture_type.name
+
+    picture_types = [
+        "NONE",
+        "I",
+        "P",
+        "B",
+        "S",
+        "SI",
+        "SP",
+        "BI",
+    ]
+
+    return picture_types[picture_type]
+
+
 class PyAVPlugin(PluginV3):
     """Support for pyAV as backend.
 
@@ -308,7 +342,7 @@ def __init__(self, request: Request, *, container: str = None, **kwargs) -> None
                     self._container = av.open(request.get_file(), **kwargs)
                 self._video_stream = self._container.streams.video[0]
                 self._decoder = self._container.decode(video=0)
-            except av.AVError:
+            except av.FFmpegError:
                 if isinstance(request.raw_uri, bytes):
                     msg = "PyAV does not support these `<bytes>`"
                 else:
@@ -458,13 +492,16 @@ def read(
 
             # reset stream container, because threading model can't change after
             # first access
-            self._video_stream.close()
             self._video_stream = self._container.streams.video[0]
 
             return frames
 
-        if thread_type is not None and thread_type != self._video_stream.thread_type:
+        if thread_type is not None and not (
+            self._video_stream.thread_type == thread_type
+            or self._video_stream.thread_type.name == thread_type
+        ):
             self._video_stream.thread_type = thread_type
+
         if (
             thread_count != 0
             and thread_count != self._video_stream.codec_context.thread_count
@@ -474,7 +511,10 @@ def read(
             self._video_stream.codec_context.thread_count = thread_count
 
         if constant_framerate is None:
-            constant_framerate = not self._container.format.variable_fps
+            # "variable_fps" is now a flag (handle got removed). Full list at
+            # https://pyav.org/docs/stable/api/container.html#module-av.format
+            variable_fps = bool(self._container.format.flags & 0x400)
+            constant_framerate = not variable_fps
 
         # note: cheap for contigous incremental reads
         self._seek(index, constant_framerate=constant_framerate)
@@ -750,7 +790,10 @@ def metadata(
             return metadata
 
         if constant_framerate is None:
-            constant_framerate = not self._container.format.variable_fps
+            # "variable_fps" is now a flag (handle got removed). Full list at
+            # https://pyav.org/docs/stable/api/container.html#module-av.format
+            variable_fps = bool(self._container.format.flags & 0x400)
+            constant_framerate = not variable_fps
 
         self._seek(index, constant_framerate=constant_framerate)
         desired_frame = next(self._decoder)
@@ -762,7 +805,7 @@ def metadata(
                 "key_frame": bool(desired_frame.key_frame),
                 "time": desired_frame.time,
                 "interlaced_frame": bool(desired_frame.interlaced_frame),
-                "frame_type": desired_frame.pict_type.name,
+                "frame_type": _get_frame_type(desired_frame.pict_type),
             }
         )
 
@@ -781,10 +824,7 @@ def close(self) -> None:
             self._flush_writer()
 
         if self._video_stream is not None:
-            try:
-                self._video_stream.close()
-            except ValueError:
-                pass  # stream already closed
+            self._video_stream = None
 
         if self._container is not None:
             self._container.close()
@@ -815,7 +855,7 @@ def init_video_stream(
         Parameters
         ----------
         codec : str
-            The codec to use, e.g. ``"libx264"`` or ``"vp9"``.
+            The codec to use, e.g. ``"h264"`` or ``"vp9"``.
         fps : float
             The desired framerate of the video stream (frames per second).
         pixel_format : str
@@ -850,7 +890,10 @@ def init_video_stream(
         if max_keyframe_interval is not None:
             stream.gop_size = max_keyframe_interval
         if force_keyframes is not None:
-            stream.closed_gop = force_keyframes
+            if force_keyframes:
+                stream.codec_context.flags |= Flags.closed_gop
+            else:
+                stream.codec_context.flags &= ~Flags.closed_gop
 
         self._video_stream = stream
 
diff --git a/tests/test_pyav.py b/tests/test_pyav.py
@@ -66,10 +66,10 @@ def test_mp4_writing(tmp_path, test_images):
         frames,
         extension=".mp4",
         plugin="pyav",
-        codec="libx264",
+        codec="h264",
     )
 
-    # libx264 writing is not deterministic and RGB -> YUV is lossy
+    # h264 writing is not deterministic and RGB -> YUV is lossy
     # so I have no good ideas how to do serious assertions on the file
     assert mp4_bytes is not None
 
@@ -224,7 +224,7 @@ def test_write_bytes(test_images):
         img,
         extension=".mp4",
         plugin="pyav",
-        codec="libx264",
+        codec="h264",
     )
 
     assert img_bytes is not None
@@ -354,7 +354,7 @@ def test_multiple_writes(test_images):
             plugin="pyav",
             format="rgba",
         ):
-            file.write(frame, is_batch=False, codec="libx264", in_pixel_format="rgba")
+            file.write(frame, is_batch=False, codec="h264", in_pixel_format="rgba")
 
     actual = out_buffer.getvalue()
     assert actual is not None
@@ -509,7 +509,7 @@ def test_procedual_writing_with_filter(test_images):
 
 def test_rotation_flag_metadata(test_images, tmp_path):
     with iio.imopen(tmp_path / "test.mp4", "w", plugin="pyav") as file:
-        file.init_video_stream("libx264")
+        file.init_video_stream("h264")
         file.container_metadata["comment"] = "This video has a rotation flag."
         file.video_stream_metadata["rotate"] = "90"
 
@@ -555,7 +555,7 @@ def test_keyframe_intervals(test_images):
         )
 
         out_file.init_video_stream(
-            "libx264", max_keyframe_interval=5, force_keyframes=True
+            "h264", max_keyframe_interval=5, force_keyframes=True
         )
 
         for idx in range(50):
@@ -568,10 +568,10 @@ def test_keyframe_intervals(test_images):
     with iio.imopen(buffer, "r", plugin="pyav") as file:
         n_frames = file.properties().shape[0]
         for idx in range(1, n_frames):
-            medatada = file.metadata(index=idx)
-            if medatada["frame_type"] == "I":
+            metatada = file.metadata(index=idx)
+            if metatada["frame_type"] == "I":
                 i_dist.append(idx)
-            if medatada["key_frame"]:
+            if metatada["key_frame"]:
                 key_dist.append(idx)
 
     assert np.max(np.diff(i_dist)) <= 5
