Fixed HWAccel so we don't share contexts between streams

Author: matthewlai

av/codec/hwaccel.pyx

@@ -1,3 +1,4 @@
+import copy
 import weakref
 from enum import IntEnum
 
@@ -7,6 +8,7 @@ from av.codec.codec cimport Codec
 from av.dictionary cimport _Dictionary
 from av.error cimport err_check
 from av.video.format cimport get_video_format
+
 from av.dictionary import Dictionary
 
 
@@ -94,34 +96,32 @@ cpdef hwdevices_available():
 
 
 cdef class HWAccel:
-    def __init__(self, device_type, device=None, codec=None, allow_software_fallback=True, options=None):
+    def __init__(self, device_type, device=None, allow_software_fallback=True, options=None):
         if isinstance(device_type, HWDeviceType):
             self._device_type = device_type
         elif isinstance(device_type, str):
             self._device_type = int(lib.av_hwdevice_find_type_by_name(device_type))
+        elif isinstance(device_type, int):
+            self._device_type = device_type
         else:
             raise ValueError("Unknown type for device_type")
 
         self._device = device
         self.allow_software_fallback = allow_software_fallback
         self.options = {} if not options else dict(options)
         self.ptr = NULL
-        self.codec = codec
         self.config = None
 
-        if codec:
-            self._initialize_hw_context()
-
-    def _initialize_hw_context(self):
+    def _initialize_hw_context(self, Codec codec not None):
         cdef HWConfig config
-        for config in self.codec.hardware_configs:
+        for config in codec.hardware_configs:
             if not (config.ptr.methods & lib.AV_CODEC_HW_CONFIG_METHOD_HW_DEVICE_CTX):
                 continue
             if self._device_type and config.device_type != self._device_type:
                 continue
             break
         else:
-            raise NotImplementedError(f"No supported hardware config for {self.codec}")
+            raise NotImplementedError(f"No supported hardware config for {codec}")
 
         self.config = config
 
@@ -142,9 +142,14 @@ cdef class HWAccel:
         if self.ptr:
             raise RuntimeError("Hardware context already initialized")
 
-        self.codec = codec
-        self._initialize_hw_context()
-        return self
+        ret = HWAccel(
+            device_type=self._device_type,
+            device=self._device,
+            allow_software_fallback=self.allow_software_fallback,
+            options=self.options
+        )
+        ret._initialize_hw_context(codec)
+        return ret
 
     def __dealloc__(self):
         if self.ptr:

examples/basics/hw_decode.py

@@ -30,11 +30,14 @@
     )
 
 if HW_DEVICE is None:
-    av.codec.hwaccel.dump_hwdevices()
-    print("Please set HW_DEVICE.")
+    print(
+        f"Please set HW_DEVICE. Options are: {str(av.codec.hwaccel.hwdevices_available())}"
+    )
     exit()
 
-assert HW_DEVICE in av.codec.hwaccel.hwdevices_available, f"{HW_DEVICE} not available."
+assert (
+    HW_DEVICE in av.codec.hwaccel.hwdevices_available()
+), f"{HW_DEVICE} not available."
 
 print("Decoding in software (auto threading)...")
 

tests/test_decode.py

@@ -8,12 +8,16 @@
 
 import av
 
+# This import is needed to make the test_decoded_time_base test pass when we run only this test file.
+# Not sure why.
+from av.subtitles import subtitle
+
 from .common import TestCase, fate_suite
 
 
 @functools.cache
 def make_h264_test_video(path: str) -> None:
-    """Generates a black H264 test video for testing hardware decoding."""
+    """Generates a black H264 test video with two streams for testing hardware decoding."""
 
     # We generate a file here that's designed to be as compatible as possible with hardware
     # encoders. Hardware encoders are sometimes very picky and the errors we get are often
@@ -23,21 +27,27 @@ def make_h264_test_video(path: str) -> None:
     # 8-bit yuv420p.
     pathlib.Path(path).parent.mkdir(parents=True, exist_ok=True)
     output_container = av.open(path, "w")
-    stream = output_container.add_stream("libx264", rate=24)
-    assert isinstance(stream, av.VideoStream)
-    stream.width = 1280
-    stream.height = 720
-    stream.pix_fmt = "yuv420p"
+
+    streams = []
+    for _ in range(2):
+        stream = output_container.add_stream("libx264", rate=24)
+        assert isinstance(stream, av.VideoStream)
+        stream.width = 1280
+        stream.height = 720
+        stream.pix_fmt = "yuv420p"
+        streams.append(stream)
 
     for _ in range(24):
         frame = av.VideoFrame.from_ndarray(
             np.zeros((720, 1280, 3), dtype=np.uint8), format="rgb24"
         )
-        for packet in stream.encode(frame):
-            output_container.mux(packet)
+        for stream in streams:
+            for packet in stream.encode(frame):
+                output_container.mux(packet)
 
-    for packet in stream.encode():
-        output_container.mux(packet)
+    for stream in streams:
+        for packet in stream.encode():
+            output_container.mux(packet)
 
     output_container.close()