Skip to content

Commit

Permalink
Fixed HWAccel so we don't share contexts between streams
Browse files Browse the repository at this point in the history
  • Loading branch information
matthewlai committed Dec 23, 2024
1 parent a4854a3 commit c10fc65
Show file tree
Hide file tree
Showing 3 changed files with 42 additions and 24 deletions.
27 changes: 16 additions & 11 deletions av/codec/hwaccel.pyx
Original file line number Diff line number Diff line change
@@ -1,3 +1,4 @@
import copy
import weakref
from enum import IntEnum

Expand All @@ -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


Expand Down Expand Up @@ -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

Expand All @@ -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:
Expand Down
9 changes: 6 additions & 3 deletions examples/basics/hw_decode.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)...")

Expand Down
30 changes: 20 additions & 10 deletions tests/test_decode.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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()

Expand Down

0 comments on commit c10fc65

Please sign in to comment.