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 authored Dec 24, 2024
1 parent a4854a3 commit af30b3d
Show file tree
Hide file tree
Showing 5 changed files with 47 additions and 29 deletions.
26 changes: 15 additions & 11 deletions av/codec/hwaccel.pyx
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,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 +95,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 +141,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
7 changes: 7 additions & 0 deletions av/container/input.pyx
Original file line number Diff line number Diff line change
Expand Up @@ -68,6 +68,8 @@ cdef class InputContainer(Container):
lib.av_dict_free(&c_options[i])
free(c_options)

at_least_one_accelerated_context = False

self.streams = StreamContainer()
for i in range(self.ptr.nb_streams):
stream = self.ptr.streams[i]
Expand All @@ -78,11 +80,16 @@ cdef class InputContainer(Container):
err_check(lib.avcodec_parameters_to_context(codec_context, stream.codecpar))
codec_context.pkt_timebase = stream.time_base
py_codec_context = wrap_codec_context(codec_context, codec, self.hwaccel)
if py_codec_context.is_hwaccel:
at_least_one_accelerated_context = True
else:
# no decoder is available
py_codec_context = None
self.streams.add_stream(wrap_stream(self, stream, py_codec_context))

if self.hwaccel and not self.hwaccel.allow_software_fallback and not at_least_one_accelerated_context:
raise RuntimeError("Hardware accelerated decode requested but no stream is compatible")

self.metadata = avdict_to_dict(self.ptr.metadata, self.metadata_encoding, self.metadata_errors)

def __dealloc__(self):
Expand Down
2 changes: 1 addition & 1 deletion av/video/codeccontext.pyx
Original file line number Diff line number Diff line change
Expand Up @@ -54,7 +54,7 @@ cdef class VideoCodecContext(CodecContext):
# stream with it, so we shouldn't abort even if we find a stream that can't
# be HW decoded.
# If the user wants to make sure hwaccel is actually used, they can check with the
# is_hardware_accelerated() function on each stream's codec context.
# is_hwaccel() function on each stream's codec context.
self.hwaccel_ctx = None

self._build_format()
Expand Down
15 changes: 8 additions & 7 deletions examples/basics/hw_decode.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@

import av
import av.datasets
from av.codec.hwaccel import HWAccel, hwdevices_available

# What accelerator to use.
# Recommendations:
Expand Down Expand Up @@ -30,11 +31,10 @@
)

if HW_DEVICE is None:
av.codec.hwaccel.dump_hwdevices()
print("Please set HW_DEVICE.")
print(f"Please set HW_DEVICE. Options are: {hwdevices_available()}")
exit()

assert HW_DEVICE in av.codec.hwaccel.hwdevices_available, f"{HW_DEVICE} not available."
assert HW_DEVICE in hwdevices_available(), f"{HW_DEVICE} not available."

print("Decoding in software (auto threading)...")

Expand All @@ -53,11 +53,12 @@
assert frame_count == container.streams.video[0].frames
container.close()

print(f"Decoded with software in {sw_time:.2f}s ({sw_fps:.2f} fps).")
print(
f"Decoded with software in {sw_time:.2f}s ({sw_fps:.2f} fps).\n"
f"Decoding with {HW_DEVICE}"
)

print(f"Decoding with {HW_DEVICE}")

hwaccel = av.codec.hwaccel.HWAccel(device_type=HW_DEVICE, allow_software_fallback=False)
hwaccel = HWAccel(device_type=HW_DEVICE, allow_software_fallback=False)

# Note the additional argument here.
container = av.open(test_file_path, hwaccel=hwaccel)
Expand Down
26 changes: 16 additions & 10 deletions tests/test_decode.py
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,7 @@

@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 +23,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 af30b3d

Please sign in to comment.