diff --git a/av/video/codeccontext.py b/av/video/codeccontext.py index eb1c9a365..fab4774a3 100644 --- a/av/video/codeccontext.py +++ b/av/video/codeccontext.py @@ -125,8 +125,7 @@ def _transfer_hwframe(self, frame: Frame): frame_sw: Frame = self._alloc_next_frame() err_check(lib.av_hwframe_transfer_data(frame_sw.ptr, frame.ptr, 0)) - # TODO: Is there anything else to transfer? - frame_sw.pts = frame.pts + frame_sw._copy_internal_attributes(frame, data_layout=False) return frame_sw @property diff --git a/av/video/reformatter.py b/av/video/reformatter.py index fbbef01b9..435038dea 100644 --- a/av/video/reformatter.py +++ b/av/video/reformatter.py @@ -287,8 +287,7 @@ def _reformat( if frame.ptr.hw_frames_ctx: frame_sw = alloc_video_frame() err_check(lib.av_hwframe_transfer_data(frame_sw.ptr, frame.ptr, 0)) - frame_sw.pts = frame.pts - frame_sw._init_user_attributes() + frame_sw._copy_internal_attributes(frame, data_layout=False) frame = frame_sw src_format = cython.cast(lib.AVPixelFormat, frame.ptr.format) diff --git a/tests/test_decode.py b/tests/test_decode.py index 3277dd776..ee4142cc4 100644 --- a/tests/test_decode.py +++ b/tests/test_decode.py @@ -272,3 +272,57 @@ def test_hardware_decode(self) -> None: frame_count += 1 assert frame_count == video_stream.frames + + +@pytest.mark.parametrize("is_hw_owned", [False, True]) +def test_hardware_decode_download_preserves_frame_props(is_hw_owned: bool) -> None: + hwdevices_available = av.codec.hwaccel.hwdevices_available() + if "HWACCEL_DEVICE_TYPE" not in os.environ: + pytest.skip( + "Set the HWACCEL_DEVICE_TYPE to run this test. " + f"Options are {' '.join(hwdevices_available)}" + ) + + hwaccel_device_type = os.environ["HWACCEL_DEVICE_TYPE"] + assert hwaccel_device_type in hwdevices_available, ( + f"{hwaccel_device_type} not available" + ) + + test_video_path = fate_suite("hevc/hdr10_plus_h265_sample.hevc") + cpu_frame = decode_first_video_frame(test_video_path) + hw_frame = decode_first_video_frame( + test_video_path, + av.codec.hwaccel.HWAccel( + device_type=hwaccel_device_type, + is_hw_owned=is_hw_owned, + allow_software_fallback=False, + ), + ) + + # Ensure that hardware decoding preserves the frame properties, see #2231 + assert_video_frame_color_props_match(hw_frame, cpu_frame) + + # Ensure that reformatting also preserves them, even for hardware frames + cpu_frame = cpu_frame.reformat(format="bgr24") + hw_frame = hw_frame.reformat(format="bgr24") + assert_video_frame_color_props_match(hw_frame, cpu_frame) + + +def decode_first_video_frame( + path: str, hwaccel: av.codec.hwaccel.HWAccel | None = None +) -> av.VideoFrame: + with av.open(path, hwaccel=hwaccel) as container: + for packet in container.demux(video=0): + frames = packet.decode() + if frames: + return frames[0] + raise AssertionError("expected at least one decoded frame") + + +def assert_video_frame_color_props_match( + actual: av.VideoFrame, expected: av.VideoFrame +) -> None: + assert actual.color_range == expected.color_range + assert actual.colorspace == expected.colorspace + assert actual.color_primaries == expected.color_primaries + assert actual.color_trc == expected.color_trc