From b35ea65277195071bc97244ef674c584a84ec2c8 Mon Sep 17 00:00:00 2001 From: WyattBlue Date: Sat, 28 Feb 2026 00:20:59 -0500 Subject: [PATCH] Add ColorTrc, ColorPrimaries, etc., closes #1968 - Add ColorTrc and ColorPrimaries IntEnum classes to av.video.reformatter - Add color_trc and color_primaries as read/write properties on VideoFrame - Add dst_color_trc and dst_color_primaries parameters to `VideoFrame.reformat()` Co-Authored-By: Claude Sonnet 4.6 --- CHANGELOG.rst | 1 + av/video/frame.py | 26 ++++++++++ av/video/frame.pyi | 5 ++ av/video/reformatter.pxd | 4 +- av/video/reformatter.py | 101 ++++++++++++++++++++++++++++++++++----- av/video/reformatter.pyi | 35 ++++++++++++++ include/avutil.pxd | 23 +++++++-- tests/test_colorspace.py | 65 ++++++++++++++++++++++++- 8 files changed, 243 insertions(+), 17 deletions(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 404b819d9..b8f944112 100644 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -41,6 +41,7 @@ Features: - Expose AVIndexEntry by :gh-user:`Queuecumber` in (:pr:`2136`). - Preserving hardware memory during cuvid decoding, exporting/importing via dlpack by :gh-user:`WyattBlue` in (:pr:`2155`). - Add enumerate_input_devices and enumerate_output_devices API by :gh-user:`WyattBlue` in (:pr:`2174`). +- Add ``ColorTrc`` and ``ColorPrimaries`` enums; add ``color_trc`` and ``color_primaries`` properties to ``VideoFrame``; add ``dst_color_trc`` and ``dst_color_primaries`` parameters to ``VideoFrame.reformat()``, addressing :issue:`1968` by :gh-user:`WyattBlue` in (:pr:`2175`). Fixes: diff --git a/av/video/frame.py b/av/video/frame.py index f6ae147cc..586d317b5 100644 --- a/av/video/frame.py +++ b/av/video/frame.py @@ -595,6 +595,32 @@ def color_range(self): def color_range(self, value): self.ptr.color_range = value + @property + def color_trc(self): + """Transfer characteristic of frame. + + Wraps :ffmpeg:`AVFrame.color_trc`. + + """ + return self.ptr.color_trc + + @color_trc.setter + def color_trc(self, value): + self.ptr.color_trc = value + + @property + def color_primaries(self): + """Color primaries of frame. + + Wraps :ffmpeg:`AVFrame.color_primaries`. + + """ + return self.ptr.color_primaries + + @color_primaries.setter + def color_primaries(self, value): + self.ptr.color_primaries = value + def reformat(self, *args, **kwargs): """reformat(width=None, height=None, format=None, src_colorspace=None, dst_colorspace=None, interpolation=None) diff --git a/av/video/frame.pyi b/av/video/frame.pyi index 570e227b0..ba3c92fed 100644 --- a/av/video/frame.pyi +++ b/av/video/frame.pyi @@ -8,6 +8,7 @@ from av.frame import Frame from .format import VideoFormat from .plane import VideoPlane +from .reformatter import ColorPrimaries, ColorTrc _SupportedNDarray = Union[ np.ndarray[Any, np.dtype[np.uint8]], @@ -41,6 +42,8 @@ class VideoFrame(Frame): pict_type: int colorspace: int color_range: int + color_trc: int + color_primaries: int @property def time(self) -> float: ... @@ -65,6 +68,8 @@ class VideoFrame(Frame): interpolation: int | str | None = None, src_color_range: int | str | None = None, dst_color_range: int | str | None = None, + dst_color_trc: int | ColorTrc | None = None, + dst_color_primaries: int | ColorPrimaries | None = None, ) -> VideoFrame: ... def to_rgb(self, **kwargs: Any) -> VideoFrame: ... def save(self, filepath: str | Path) -> None: ... diff --git a/av/video/reformatter.pxd b/av/video/reformatter.pxd index 433824e8c..f57ba7e77 100644 --- a/av/video/reformatter.pxd +++ b/av/video/reformatter.pxd @@ -79,4 +79,6 @@ cdef class VideoReformatter: lib.AVPixelFormat format, int src_colorspace, int dst_colorspace, int interpolation, int src_color_range, int dst_color_range, - bint set_dst_colorspace, bint set_dst_color_range) + bint set_dst_colorspace, bint set_dst_color_range, + int dst_color_trc, int dst_color_primaries, + bint set_dst_color_trc, bint set_dst_color_primaries) diff --git a/av/video/reformatter.py b/av/video/reformatter.py index b72162ede..9c201bbe9 100644 --- a/av/video/reformatter.py +++ b/av/video/reformatter.py @@ -44,7 +44,55 @@ class ColorRange(IntEnum): NB: "Not part of ABI" = lib.AVCOL_RANGE_NB -def _resolve_enum_value(value, enum_class, default): +class ColorTrc(IntEnum): + """Transfer characteristic (gamma curve) of a video frame. + + Maps to FFmpeg's ``AVColorTransferCharacteristic``. + """ + + BT709: "BT.709" = lib.AVCOL_TRC_BT709 + UNSPECIFIED: "Unspecified" = lib.AVCOL_TRC_UNSPECIFIED + GAMMA22: "Gamma 2.2 (BT.470M)" = lib.AVCOL_TRC_GAMMA22 + GAMMA28: "Gamma 2.8 (BT.470BG)" = lib.AVCOL_TRC_GAMMA28 + SMPTE170M: "SMPTE 170M" = lib.AVCOL_TRC_SMPTE170M + SMPTE240M: "SMPTE 240M" = lib.AVCOL_TRC_SMPTE240M + LINEAR: "Linear" = lib.AVCOL_TRC_LINEAR + LOG: "Logarithmic (100:1 range)" = lib.AVCOL_TRC_LOG + LOG_SQRT: "Logarithmic (100*sqrt(10):1 range)" = lib.AVCOL_TRC_LOG_SQRT + IEC61966_2_4: "IEC 61966-2-4 (sRGB)" = lib.AVCOL_TRC_IEC61966_2_4 + BT1361_ECG: "BT.1361 extended colour gamut" = lib.AVCOL_TRC_BT1361_ECG + IEC61966_2_1: "IEC 61966-2-1 (sYCC)" = lib.AVCOL_TRC_IEC61966_2_1 + BT2020_10: "BT.2020 10-bit" = lib.AVCOL_TRC_BT2020_10 + BT2020_12: "BT.2020 12-bit" = lib.AVCOL_TRC_BT2020_12 + SMPTE2084: "SMPTE 2084 (PQ, HDR10)" = lib.AVCOL_TRC_SMPTE2084 + SMPTE428: "SMPTE 428-1" = lib.AVCOL_TRC_SMPTE428 + ARIB_STD_B67: "ARIB STD-B67 (HLG)" = lib.AVCOL_TRC_ARIB_STD_B67 + + +class ColorPrimaries(IntEnum): + """Color primaries of a video frame. + + Maps to FFmpeg's ``AVColorPrimaries``. + """ + + BT709: "BT.709 / sRGB / sYCC" = lib.AVCOL_PRI_BT709 + UNSPECIFIED: "Unspecified" = lib.AVCOL_PRI_UNSPECIFIED + BT470M: "BT.470M" = lib.AVCOL_PRI_BT470M + BT470BG: "BT.470BG / BT.601-6 625" = lib.AVCOL_PRI_BT470BG + SMPTE170M: "SMPTE 170M / BT.601-6 525" = lib.AVCOL_PRI_SMPTE170M + SMPTE240M: "SMPTE 240M" = lib.AVCOL_PRI_SMPTE240M + FILM: "Generic film (Illuminant C)" = lib.AVCOL_PRI_FILM + BT2020: "BT.2020 / BT.2100" = lib.AVCOL_PRI_BT2020 + SMPTE428: "SMPTE 428-1 / XYZ" = lib.AVCOL_PRI_SMPTE428 + SMPTE431: "SMPTE 431-2 (DCI-P3)" = lib.AVCOL_PRI_SMPTE431 + SMPTE432: "SMPTE 432-1 (Display P3)" = lib.AVCOL_PRI_SMPTE432 + EBU3213: "EBU 3213-E / JEDEC P22" = lib.AVCOL_PRI_EBU3213 + + +@cython.cfunc +def _resolve_enum_value( + value: object, enum_class: object, default: cython.int +) -> cython.int: # Helper function to resolve enum values from different input types. if value is None: return default @@ -96,6 +144,8 @@ def reformat( interpolation=None, src_color_range=None, dst_color_range=None, + dst_color_trc=None, + dst_color_primaries=None, ): """Create a new :class:`VideoFrame` with the given width/height/format/colorspace. @@ -112,34 +162,43 @@ def reformat( :param interpolation: The interpolation method to use, or ``None`` for ``BILINEAR``. :type interpolation: :class:`Interpolation` or ``str`` :param src_color_range: Current color range, or ``None`` for the ``UNSPECIFIED``. - :type src_color_range: :class:`color range` or ``str`` + :type src_color_range: :class:`ColorRange` or ``str`` :param dst_color_range: Desired color range, or ``None`` for the ``UNSPECIFIED``. - :type dst_color_range: :class:`color range` or ``str`` + :type dst_color_range: :class:`ColorRange` or ``str`` + :param dst_color_trc: Desired transfer characteristic to tag on the output frame, + or ``None`` to preserve the source frame's value. This sets frame metadata only; + it does not perform a pixel-level transfer function conversion. + :type dst_color_trc: :class:`ColorTrc` or ``int`` + :param dst_color_primaries: Desired color primaries to tag on the output frame, + or ``None`` to preserve the source frame's value. + :type dst_color_primaries: :class:`ColorPrimaries` or ``int`` """ video_format: VideoFormat = VideoFormat( format if format is not None else frame.format ) - c_src_colorspace: cython.int = _resolve_enum_value( + c_src_colorspace = _resolve_enum_value( src_colorspace, Colorspace, frame.colorspace ) - c_dst_colorspace: cython.int = _resolve_enum_value( + c_dst_colorspace = _resolve_enum_value( dst_colorspace, Colorspace, frame.colorspace ) - c_interpolation: cython.int = _resolve_enum_value( + c_interpolation = _resolve_enum_value( interpolation, Interpolation, int(Interpolation.BILINEAR) ) - c_src_color_range: cython.int = _resolve_enum_value( - src_color_range, ColorRange, 0 - ) - c_dst_color_range: cython.int = _resolve_enum_value( - dst_color_range, ColorRange, 0 + c_src_color_range = _resolve_enum_value(src_color_range, ColorRange, 0) + c_dst_color_range = _resolve_enum_value(dst_color_range, ColorRange, 0) + c_dst_color_trc = _resolve_enum_value(dst_color_trc, ColorTrc, 0) + c_dst_color_primaries = _resolve_enum_value( + dst_color_primaries, ColorPrimaries, 0 ) # Track whether user explicitly specified destination metadata set_dst_colorspace: cython.bint = dst_colorspace is not None set_dst_color_range: cython.bint = dst_color_range is not None + set_dst_color_trc: cython.bint = dst_color_trc is not None + set_dst_color_primaries: cython.bint = dst_color_primaries is not None return self._reformat( frame, @@ -153,6 +212,10 @@ def reformat( c_dst_color_range, set_dst_colorspace, set_dst_color_range, + c_dst_color_trc, + c_dst_color_primaries, + set_dst_color_trc, + set_dst_color_primaries, ) @cython.cfunc @@ -169,6 +232,10 @@ def _reformat( dst_color_range: cython.int, set_dst_colorspace: cython.bint, set_dst_color_range: cython.bint, + dst_color_trc: cython.int, + dst_color_primaries: cython.int, + set_dst_color_trc: cython.bint, + set_dst_color_primaries: cython.bint, ): if frame.ptr.format < 0: raise ValueError("Frame does not have format set.") @@ -191,6 +258,8 @@ def _reformat( and height == frame.ptr.height and dst_colorspace == src_colorspace and src_color_range == dst_color_range + and not set_dst_color_trc + and not set_dst_color_primaries ): return frame @@ -207,6 +276,8 @@ def _reformat( and height == frame.ptr.height and dst_colorspace == src_colorspace and src_color_range == dst_color_range + and not set_dst_color_trc + and not set_dst_color_primaries ): return frame @@ -285,6 +356,14 @@ def _reformat( new_frame.ptr.color_range = cython.cast( lib.AVColorRange, frame_dst_color_range ) + if set_dst_color_trc: + new_frame.ptr.color_trc = cython.cast( + lib.AVColorTransferCharacteristic, dst_color_trc + ) + if set_dst_color_primaries: + new_frame.ptr.color_primaries = cython.cast( + lib.AVColorPrimaries, dst_color_primaries + ) with cython.nogil: sws_scale( diff --git a/av/video/reformatter.pyi b/av/video/reformatter.pyi index 5d83fcbe3..b1e51e984 100644 --- a/av/video/reformatter.pyi +++ b/av/video/reformatter.pyi @@ -38,6 +38,39 @@ class ColorRange(IntEnum): JPEG = 2 NB = 3 +class ColorTrc(IntEnum): + BT709 = cast(int, ...) + UNSPECIFIED = cast(int, ...) + GAMMA22 = cast(int, ...) + GAMMA28 = cast(int, ...) + SMPTE170M = cast(int, ...) + SMPTE240M = cast(int, ...) + LINEAR = cast(int, ...) + LOG = cast(int, ...) + LOG_SQRT = cast(int, ...) + IEC61966_2_4 = cast(int, ...) + BT1361_ECG = cast(int, ...) + IEC61966_2_1 = cast(int, ...) + BT2020_10 = cast(int, ...) + BT2020_12 = cast(int, ...) + SMPTE2084 = cast(int, ...) + SMPTE428 = cast(int, ...) + ARIB_STD_B67 = cast(int, ...) + +class ColorPrimaries(IntEnum): + BT709 = cast(int, ...) + UNSPECIFIED = cast(int, ...) + BT470M = cast(int, ...) + BT470BG = cast(int, ...) + SMPTE170M = cast(int, ...) + SMPTE240M = cast(int, ...) + FILM = cast(int, ...) + BT2020 = cast(int, ...) + SMPTE428 = cast(int, ...) + SMPTE431 = cast(int, ...) + SMPTE432 = cast(int, ...) + EBU3213 = cast(int, ...) + class VideoReformatter: def reformat( self, @@ -50,4 +83,6 @@ class VideoReformatter: interpolation: int | str | None = None, src_color_range: int | str | None = None, dst_color_range: int | str | None = None, + dst_color_trc: int | ColorTrc | None = None, + dst_color_primaries: int | ColorPrimaries | None = None, ) -> VideoFrame: ... diff --git a/include/avutil.pxd b/include/avutil.pxd index c1320a97d..3ee78d2f5 100644 --- a/include/avutil.pxd +++ b/include/avutil.pxd @@ -53,10 +53,8 @@ cdef extern from "libavutil/avutil.h" nogil: AVCOL_RANGE_NB cdef enum AVColorPrimaries: - AVCOL_PRI_RESERVED0 AVCOL_PRI_BT709 AVCOL_PRI_UNSPECIFIED - AVCOL_PRI_RESERVED AVCOL_PRI_BT470M AVCOL_PRI_BT470BG AVCOL_PRI_SMPTE170M @@ -69,10 +67,27 @@ cdef extern from "libavutil/avutil.h" nogil: AVCOL_PRI_SMPTE432 AVCOL_PRI_EBU3213 AVCOL_PRI_JEDEC_P22 - AVCOL_PRI_NB cdef enum AVColorTransferCharacteristic: - pass + AVCOL_TRC_BT709 + AVCOL_TRC_UNSPECIFIED + AVCOL_TRC_GAMMA22 + AVCOL_TRC_GAMMA28 + AVCOL_TRC_SMPTE170M + AVCOL_TRC_SMPTE240M + AVCOL_TRC_LINEAR + AVCOL_TRC_LOG + AVCOL_TRC_LOG_SQRT + AVCOL_TRC_IEC61966_2_4 + AVCOL_TRC_BT1361_ECG + AVCOL_TRC_IEC61966_2_1 + AVCOL_TRC_BT2020_10 + AVCOL_TRC_BT2020_12 + AVCOL_TRC_SMPTE2084 + AVCOL_TRC_SMPTEST2084 + AVCOL_TRC_SMPTE428 + AVCOL_TRC_SMPTEST428_1 + AVCOL_TRC_ARIB_STD_B67 cdef void* av_malloc(size_t size) cdef void* av_mallocz(size_t size) diff --git a/tests/test_colorspace.py b/tests/test_colorspace.py index c76416c80..a5352fc7d 100644 --- a/tests/test_colorspace.py +++ b/tests/test_colorspace.py @@ -1,5 +1,10 @@ import av -from av.video.reformatter import ColorRange, Colorspace +from av.video.reformatter import ( + ColorPrimaries, + ColorRange, + Colorspace, + ColorTrc, +) from .common import fate_suite @@ -38,3 +43,61 @@ def test_sky_timelapse() -> None: assert stream.codec_context.color_primaries == 1 assert stream.codec_context.color_trc == 1 assert stream.codec_context.colorspace == 1 + + +def test_frame_color_trc_property() -> None: + frame = av.VideoFrame(width=64, height=64, format="rgb24") + assert frame.color_trc == ColorTrc.UNSPECIFIED + + frame.color_trc = ColorTrc.IEC61966_2_4 + assert frame.color_trc == ColorTrc.IEC61966_2_4 + + frame.color_trc = ColorTrc.BT709 + assert frame.color_trc == ColorTrc.BT709 + + +def test_frame_color_primaries_property() -> None: + frame = av.VideoFrame(width=64, height=64, format="rgb24") + assert frame.color_primaries == ColorPrimaries.UNSPECIFIED + + frame.color_primaries = ColorPrimaries.BT709 + assert frame.color_primaries == ColorPrimaries.BT709 + assert frame.color_primaries == 1 # AVCOL_PRI_BT709 + + +def test_reformat_dst_color_trc() -> None: + # Reformat a frame and tag it with sRGB transfer characteristic. + frame = av.VideoFrame(width=64, height=64, format="yuv420p") + rgb = frame.reformat( + format="rgb24", + dst_colorspace=Colorspace.ITU709, + dst_color_trc=ColorTrc.IEC61966_2_4, + ) + assert rgb.format.name == "rgb24" + assert rgb.colorspace == Colorspace.ITU709 + assert rgb.color_trc == ColorTrc.IEC61966_2_4 + + +def test_reformat_dst_color_primaries() -> None: + frame = av.VideoFrame(width=64, height=64, format="yuv420p") + rgb = frame.reformat( + format="rgb24", + dst_color_primaries=ColorPrimaries.BT709, + ) + assert rgb.color_primaries == ColorPrimaries.BT709 + + +def test_reformat_preserves_color_trc() -> None: + # When dst_color_trc is not specified, the source frame's value is preserved. + frame = av.VideoFrame(width=64, height=64, format="yuv420p") + frame.color_trc = ColorTrc.BT709 + rgb = frame.reformat(format="rgb24") + assert rgb.color_trc == ColorTrc.BT709 + + +def test_reformat_preserves_color_primaries() -> None: + # When dst_color_primaries is not specified, the source frame's value is preserved. + frame = av.VideoFrame(width=64, height=64, format="yuv420p") + frame.color_primaries = ColorPrimaries.BT709 + rgb = frame.reformat(format="rgb24") + assert rgb.color_primaries == ColorPrimaries.BT709