Add Colorspace.BT2020 for BT.2020 YUV matrix conversion

Author: mark-oshea

CHANGELOG.rst

@@ -41,6 +41,7 @@ Features:
 - Support ``yuv420p10le`` in ``VideoFrame.to_ndarray`` and ``VideoFrame.from_ndarray`` by :gh-user:`WyattBlue` (:issue:`1981`).
 - Add ``at`` parameter to ``Graph.push`` and ``Graph.vpush`` to push a frame to a single buffer source by index, for multi-input filters like ``overlay`` by :gh-user:`WyattBlue`.
 - ``find_best_pix_fmt_of_list`` now returns the loss as a ``PixFmtLoss`` ``enum.IntFlag`` instead of a plain ``int`` by :gh-user:`WyattBlue` (:issue:`2300`).
+- Add ``Colorspace.BT2020`` by :gh-user:`mark-oshea`.
 
 Fixes:
 

av/video/reformatter.pxd

@@ -33,6 +33,7 @@ cdef extern from "libswscale/swscale.h" nogil:
     cdef int SWS_CS_SMPTE170M
     cdef int SWS_CS_SMPTE240M
     cdef int SWS_CS_DEFAULT
+    cdef int SWS_CS_BT2020
 
     cdef SwsContext *sws_alloc_context()
     cdef void sws_free_context(SwsContext **ctx)

av/video/reformatter.py

@@ -36,6 +36,7 @@ class Colorspace(IntEnum):
     SMPTE170M = SWS_CS_SMPTE170M
     SMPTE240M = SWS_CS_SMPTE240M
     DEFAULT = SWS_CS_DEFAULT
+    BT2020 = SWS_CS_BT2020
     # Lowercase for b/c.
     itu709 = SWS_CS_ITU709
     fcc = SWS_CS_FCC
@@ -44,6 +45,7 @@ class Colorspace(IntEnum):
     smpte170m = SWS_CS_SMPTE170M
     smpte240m = SWS_CS_SMPTE240M
     default = SWS_CS_DEFAULT
+    bt2020 = SWS_CS_BT2020
 
 
 class ColorRange(IntEnum):
@@ -146,6 +148,8 @@ def _set_frame_colorspace(
         frame.colorspace = lib.AVCOL_SPC_SMPTE170M
     elif colorspace == SWS_CS_SMPTE240M:
         frame.colorspace = lib.AVCOL_SPC_SMPTE240M
+    elif colorspace == SWS_CS_BT2020:
+        frame.colorspace = lib.AVCOL_SPC_BT2020_NCL
 
 
 @cython.final

av/video/reformatter.pyi

@@ -31,13 +31,15 @@ class Colorspace(IntEnum):
     SMPTE170M = cast(int, ...)
     SMPTE240M = cast(int, ...)
     DEFAULT = cast(int, ...)
+    BT2020 = cast(int, ...)
     itu709 = cast(int, ...)
     fcc = cast(int, ...)
     itu601 = cast(int, ...)
     itu624 = cast(int, ...)
     smpte170m = cast(int, ...)
     smpte240m = cast(int, ...)
     default = cast(int, ...)
+    bt2020 = cast(int, ...)
 
 class ColorRange(IntEnum):
     UNSPECIFIED = 0

tests/test_colorspace.py

@@ -1,4 +1,5 @@
 import av
+import pytest
 from av.video.reformatter import (
     ColorPrimaries,
     ColorRange,
@@ -101,3 +102,21 @@ def test_reformat_preserves_color_primaries() -> None:
     frame.color_primaries = ColorPrimaries.BT709
     rgb = frame.reformat(format="rgb24")
     assert rgb.color_primaries == ColorPrimaries.BT709
+
+
+@pytest.mark.parametrize(
+    ("colorspace", "expected"),
+    [
+        (Colorspace.ITU709, Colorspace.ITU709),
+        (Colorspace.FCC, Colorspace.FCC),
+        (Colorspace.ITU601, 6),  # AVCOL_SPC_SMPTE170M
+        (Colorspace.SMPTE240M, Colorspace.SMPTE240M),
+        (Colorspace.BT2020, Colorspace.BT2020),
+    ],
+)
+def test_reformat_dst_colorspace_metadata(
+    colorspace: Colorspace, expected: Colorspace | int
+) -> None:
+    frame = av.VideoFrame(width=64, height=64, format="yuv420p")
+    rgb = frame.reformat(format="rgb24", dst_colorspace=colorspace)
+    assert rgb.colorspace == expected