Skip to content

Commit 8760006

Browse files
WyattBlueclaude
WyattBlue
and
claude
committed
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 <noreply@anthropic.com>
1 parent 0a1ed6f commit 8760006

8 files changed

Lines changed: 243 additions & 17 deletions

File tree

‎CHANGELOG.rst‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -41,6 +41,7 @@ Features:
4141
- Expose AVIndexEntry by :gh-user:`Queuecumber` in (:pr:`2136`).
4242
- Preserving hardware memory during cuvid decoding, exporting/importing via dlpack by :gh-user:`WyattBlue` in (:pr:`2155`).
4343
- Add enumerate_input_devices and enumerate_output_devices API by :gh-user:`WyattBlue` in (:pr:`2174`).
44+
- 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`).
4445

4546
Fixes:
4647

‎av/video/frame.py‎

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -595,6 +595,32 @@ def color_range(self):
595595
def color_range(self, value):
596596
self.ptr.color_range = value
597597

598+
@property
599+
def color_trc(self):
600+
"""Transfer characteristic of frame.
601+
602+
Wraps :ffmpeg:`AVFrame.color_trc`.
603+
604+
"""
605+
return self.ptr.color_trc
606+
607+
@color_trc.setter
608+
def color_trc(self, value):
609+
self.ptr.color_trc = value
610+
611+
@property
612+
def color_primaries(self):
613+
"""Color primaries of frame.
614+
615+
Wraps :ffmpeg:`AVFrame.color_primaries`.
616+
617+
"""
618+
return self.ptr.color_primaries
619+
620+
@color_primaries.setter
621+
def color_primaries(self, value):
622+
self.ptr.color_primaries = value
623+
598624
def reformat(self, *args, **kwargs):
599625
"""reformat(width=None, height=None, format=None, src_colorspace=None, dst_colorspace=None, interpolation=None)
600626

‎av/video/frame.pyi‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@ from av.frame import Frame
88

99
from .format import VideoFormat
1010
from .plane import VideoPlane
11+
from .reformatter import ColorPrimaries, ColorTrc
1112

1213
_SupportedNDarray = Union[
1314
np.ndarray[Any, np.dtype[np.uint8]],
@@ -41,6 +42,8 @@ class VideoFrame(Frame):
4142
pict_type: int
4243
colorspace: int
4344
color_range: int
45+
color_trc: int
46+
color_primaries: int
4447

4548
@property
4649
def time(self) -> float: ...
@@ -65,6 +68,8 @@ class VideoFrame(Frame):
6568
interpolation: int | str | None = None,
6669
src_color_range: int | str | None = None,
6770
dst_color_range: int | str | None = None,
71+
dst_color_trc: int | ColorTrc | None = None,
72+
dst_color_primaries: int | ColorPrimaries | None = None,
6873
) -> VideoFrame: ...
6974
def to_rgb(self, **kwargs: Any) -> VideoFrame: ...
7075
def save(self, filepath: str | Path) -> None: ...

‎av/video/reformatter.pxd‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -79,4 +79,6 @@ cdef class VideoReformatter:
7979
lib.AVPixelFormat format, int src_colorspace,
8080
int dst_colorspace, int interpolation,
8181
int src_color_range, int dst_color_range,
82-
bint set_dst_colorspace, bint set_dst_color_range)
82+
bint set_dst_colorspace, bint set_dst_color_range,
83+
int dst_color_trc, int dst_color_primaries,
84+
bint set_dst_color_trc, bint set_dst_color_primaries)

‎av/video/reformatter.py‎

Lines changed: 90 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -44,7 +44,55 @@ class ColorRange(IntEnum):
4444
NB: "Not part of ABI" = lib.AVCOL_RANGE_NB
4545

4646

47-
def _resolve_enum_value(value, enum_class, default):
47+
class ColorTrc(IntEnum):
48+
"""Transfer characteristic (gamma curve) of a video frame.
49+
50+
Maps to FFmpeg's ``AVColorTransferCharacteristic``.
51+
"""
52+
53+
BT709: "BT.709" = lib.AVCOL_TRC_BT709
54+
UNSPECIFIED: "Unspecified" = lib.AVCOL_TRC_UNSPECIFIED
55+
GAMMA22: "Gamma 2.2 (BT.470M)" = lib.AVCOL_TRC_GAMMA22
56+
GAMMA28: "Gamma 2.8 (BT.470BG)" = lib.AVCOL_TRC_GAMMA28
57+
SMPTE170M: "SMPTE 170M" = lib.AVCOL_TRC_SMPTE170M
58+
SMPTE240M: "SMPTE 240M" = lib.AVCOL_TRC_SMPTE240M
59+
LINEAR: "Linear" = lib.AVCOL_TRC_LINEAR
60+
LOG: "Logarithmic (100:1 range)" = lib.AVCOL_TRC_LOG
61+
LOG_SQRT: "Logarithmic (100*sqrt(10):1 range)" = lib.AVCOL_TRC_LOG_SQRT
62+
IEC61966_2_4: "IEC 61966-2-4 (sRGB)" = lib.AVCOL_TRC_IEC61966_2_4
63+
BT1361_ECG: "BT.1361 extended colour gamut" = lib.AVCOL_TRC_BT1361_ECG
64+
IEC61966_2_1: "IEC 61966-2-1 (sYCC)" = lib.AVCOL_TRC_IEC61966_2_1
65+
BT2020_10: "BT.2020 10-bit" = lib.AVCOL_TRC_BT2020_10
66+
BT2020_12: "BT.2020 12-bit" = lib.AVCOL_TRC_BT2020_12
67+
SMPTE2084: "SMPTE 2084 (PQ, HDR10)" = lib.AVCOL_TRC_SMPTE2084
68+
SMPTE428: "SMPTE 428-1" = lib.AVCOL_TRC_SMPTE428
69+
ARIB_STD_B67: "ARIB STD-B67 (HLG)" = lib.AVCOL_TRC_ARIB_STD_B67
70+
71+
72+
class ColorPrimaries(IntEnum):
73+
"""Color primaries of a video frame.
74+
75+
Maps to FFmpeg's ``AVColorPrimaries``.
76+
"""
77+
78+
BT709: "BT.709 / sRGB / sYCC" = lib.AVCOL_PRI_BT709
79+
UNSPECIFIED: "Unspecified" = lib.AVCOL_PRI_UNSPECIFIED
80+
BT470M: "BT.470M" = lib.AVCOL_PRI_BT470M
81+
BT470BG: "BT.470BG / BT.601-6 625" = lib.AVCOL_PRI_BT470BG
82+
SMPTE170M: "SMPTE 170M / BT.601-6 525" = lib.AVCOL_PRI_SMPTE170M
83+
SMPTE240M: "SMPTE 240M" = lib.AVCOL_PRI_SMPTE240M
84+
FILM: "Generic film (Illuminant C)" = lib.AVCOL_PRI_FILM
85+
BT2020: "BT.2020 / BT.2100" = lib.AVCOL_PRI_BT2020
86+
SMPTE428: "SMPTE 428-1 / XYZ" = lib.AVCOL_PRI_SMPTE428
87+
SMPTE431: "SMPTE 431-2 (DCI-P3)" = lib.AVCOL_PRI_SMPTE431
88+
SMPTE432: "SMPTE 432-1 (Display P3)" = lib.AVCOL_PRI_SMPTE432
89+
EBU3213: "EBU 3213-E / JEDEC P22" = lib.AVCOL_PRI_EBU3213
90+
91+
92+
@cython.cfunc
93+
def _resolve_enum_value(
94+
value: object, enum_class: object, default: cython.int
95+
) -> cython.int:
4896
# Helper function to resolve enum values from different input types.
4997
if value is None:
5098
return default
@@ -96,6 +144,8 @@ def reformat(
96144
interpolation=None,
97145
src_color_range=None,
98146
dst_color_range=None,
147+
dst_color_trc=None,
148+
dst_color_primaries=None,
99149
):
100150
"""Create a new :class:`VideoFrame` with the given width/height/format/colorspace.
101151
@@ -112,34 +162,43 @@ def reformat(
112162
:param interpolation: The interpolation method to use, or ``None`` for ``BILINEAR``.
113163
:type interpolation: :class:`Interpolation` or ``str``
114164
:param src_color_range: Current color range, or ``None`` for the ``UNSPECIFIED``.
115-
:type src_color_range: :class:`color range` or ``str``
165+
:type src_color_range: :class:`ColorRange` or ``str``
116166
:param dst_color_range: Desired color range, or ``None`` for the ``UNSPECIFIED``.
117-
:type dst_color_range: :class:`color range` or ``str``
167+
:type dst_color_range: :class:`ColorRange` or ``str``
168+
:param dst_color_trc: Desired transfer characteristic to tag on the output frame,
169+
or ``None`` to preserve the source frame's value. This sets frame metadata only;
170+
it does not perform a pixel-level transfer function conversion.
171+
:type dst_color_trc: :class:`ColorTrc` or ``int``
172+
:param dst_color_primaries: Desired color primaries to tag on the output frame,
173+
or ``None`` to preserve the source frame's value.
174+
:type dst_color_primaries: :class:`ColorPrimaries` or ``int``
118175
119176
"""
120177

121178
video_format: VideoFormat = VideoFormat(
122179
format if format is not None else frame.format
123180
)
124-
c_src_colorspace: cython.int = _resolve_enum_value(
181+
c_src_colorspace = _resolve_enum_value(
125182
src_colorspace, Colorspace, frame.colorspace
126183
)
127-
c_dst_colorspace: cython.int = _resolve_enum_value(
184+
c_dst_colorspace = _resolve_enum_value(
128185
dst_colorspace, Colorspace, frame.colorspace
129186
)
130-
c_interpolation: cython.int = _resolve_enum_value(
187+
c_interpolation = _resolve_enum_value(
131188
interpolation, Interpolation, int(Interpolation.BILINEAR)
132189
)
133-
c_src_color_range: cython.int = _resolve_enum_value(
134-
src_color_range, ColorRange, 0
135-
)
136-
c_dst_color_range: cython.int = _resolve_enum_value(
137-
dst_color_range, ColorRange, 0
190+
c_src_color_range = _resolve_enum_value(src_color_range, ColorRange, 0)
191+
c_dst_color_range = _resolve_enum_value(dst_color_range, ColorRange, 0)
192+
c_dst_color_trc = _resolve_enum_value(dst_color_trc, ColorTrc, 0)
193+
c_dst_color_primaries = _resolve_enum_value(
194+
dst_color_primaries, ColorPrimaries, 0
138195
)
139196

140197
# Track whether user explicitly specified destination metadata
141198
set_dst_colorspace: cython.bint = dst_colorspace is not None
142199
set_dst_color_range: cython.bint = dst_color_range is not None
200+
set_dst_color_trc: cython.bint = dst_color_trc is not None
201+
set_dst_color_primaries: cython.bint = dst_color_primaries is not None
143202

144203
return self._reformat(
145204
frame,
@@ -153,6 +212,10 @@ def reformat(
153212
c_dst_color_range,
154213
set_dst_colorspace,
155214
set_dst_color_range,
215+
c_dst_color_trc,
216+
c_dst_color_primaries,
217+
set_dst_color_trc,
218+
set_dst_color_primaries,
156219
)
157220

158221
@cython.cfunc
@@ -169,6 +232,10 @@ def _reformat(
169232
dst_color_range: cython.int,
170233
set_dst_colorspace: cython.bint,
171234
set_dst_color_range: cython.bint,
235+
dst_color_trc: cython.int,
236+
dst_color_primaries: cython.int,
237+
set_dst_color_trc: cython.bint,
238+
set_dst_color_primaries: cython.bint,
172239
):
173240
if frame.ptr.format < 0:
174241
raise ValueError("Frame does not have format set.")
@@ -191,6 +258,8 @@ def _reformat(
191258
and height == frame.ptr.height
192259
and dst_colorspace == src_colorspace
193260
and src_color_range == dst_color_range
261+
and not set_dst_color_trc
262+
and not set_dst_color_primaries
194263
):
195264
return frame
196265

@@ -207,6 +276,8 @@ def _reformat(
207276
and height == frame.ptr.height
208277
and dst_colorspace == src_colorspace
209278
and src_color_range == dst_color_range
279+
and not set_dst_color_trc
280+
and not set_dst_color_primaries
210281
):
211282
return frame
212283

@@ -285,6 +356,14 @@ def _reformat(
285356
new_frame.ptr.color_range = cython.cast(
286357
lib.AVColorRange, frame_dst_color_range
287358
)
359+
if set_dst_color_trc:
360+
new_frame.ptr.color_trc = cython.cast(
361+
lib.AVColorTransferCharacteristic, dst_color_trc
362+
)
363+
if set_dst_color_primaries:
364+
new_frame.ptr.color_primaries = cython.cast(
365+
lib.AVColorPrimaries, dst_color_primaries
366+
)
288367

289368
with cython.nogil:
290369
sws_scale(

‎av/video/reformatter.pyi‎

Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -38,6 +38,39 @@ class ColorRange(IntEnum):
3838
JPEG = 2
3939
NB = 3
4040

41+
class ColorTrc(IntEnum):
42+
BT709 = cast(int, ...)
43+
UNSPECIFIED = cast(int, ...)
44+
GAMMA22 = cast(int, ...)
45+
GAMMA28 = cast(int, ...)
46+
SMPTE170M = cast(int, ...)
47+
SMPTE240M = cast(int, ...)
48+
LINEAR = cast(int, ...)
49+
LOG = cast(int, ...)
50+
LOG_SQRT = cast(int, ...)
51+
IEC61966_2_4 = cast(int, ...)
52+
BT1361_ECG = cast(int, ...)
53+
IEC61966_2_1 = cast(int, ...)
54+
BT2020_10 = cast(int, ...)
55+
BT2020_12 = cast(int, ...)
56+
SMPTE2084 = cast(int, ...)
57+
SMPTE428 = cast(int, ...)
58+
ARIB_STD_B67 = cast(int, ...)
59+
60+
class ColorPrimaries(IntEnum):
61+
BT709 = cast(int, ...)
62+
UNSPECIFIED = cast(int, ...)
63+
BT470M = cast(int, ...)
64+
BT470BG = cast(int, ...)
65+
SMPTE170M = cast(int, ...)
66+
SMPTE240M = cast(int, ...)
67+
FILM = cast(int, ...)
68+
BT2020 = cast(int, ...)
69+
SMPTE428 = cast(int, ...)
70+
SMPTE431 = cast(int, ...)
71+
SMPTE432 = cast(int, ...)
72+
EBU3213 = cast(int, ...)
73+
4174
class VideoReformatter:
4275
def reformat(
4376
self,
@@ -50,4 +83,6 @@ class VideoReformatter:
5083
interpolation: int | str | None = None,
5184
src_color_range: int | str | None = None,
5285
dst_color_range: int | str | None = None,
86+
dst_color_trc: int | ColorTrc | None = None,
87+
dst_color_primaries: int | ColorPrimaries | None = None,
5388
) -> VideoFrame: ...

‎include/avutil.pxd‎

Lines changed: 19 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -53,10 +53,8 @@ cdef extern from "libavutil/avutil.h" nogil:
5353
AVCOL_RANGE_NB
5454

5555
cdef enum AVColorPrimaries:
56-
AVCOL_PRI_RESERVED0
5756
AVCOL_PRI_BT709
5857
AVCOL_PRI_UNSPECIFIED
59-
AVCOL_PRI_RESERVED
6058
AVCOL_PRI_BT470M
6159
AVCOL_PRI_BT470BG
6260
AVCOL_PRI_SMPTE170M
@@ -69,10 +67,27 @@ cdef extern from "libavutil/avutil.h" nogil:
6967
AVCOL_PRI_SMPTE432
7068
AVCOL_PRI_EBU3213
7169
AVCOL_PRI_JEDEC_P22
72-
AVCOL_PRI_NB
7370

7471
cdef enum AVColorTransferCharacteristic:
75-
pass
72+
AVCOL_TRC_BT709
73+
AVCOL_TRC_UNSPECIFIED
74+
AVCOL_TRC_GAMMA22
75+
AVCOL_TRC_GAMMA28
76+
AVCOL_TRC_SMPTE170M
77+
AVCOL_TRC_SMPTE240M
78+
AVCOL_TRC_LINEAR
79+
AVCOL_TRC_LOG
80+
AVCOL_TRC_LOG_SQRT
81+
AVCOL_TRC_IEC61966_2_4
82+
AVCOL_TRC_BT1361_ECG
83+
AVCOL_TRC_IEC61966_2_1
84+
AVCOL_TRC_BT2020_10
85+
AVCOL_TRC_BT2020_12
86+
AVCOL_TRC_SMPTE2084
87+
AVCOL_TRC_SMPTEST2084
88+
AVCOL_TRC_SMPTE428
89+
AVCOL_TRC_SMPTEST428_1
90+
AVCOL_TRC_ARIB_STD_B67
7691

7792
cdef void* av_malloc(size_t size)
7893
cdef void* av_mallocz(size_t size)

0 commit comments

Comments
 (0)