Expose the chroma location of a video frame

CHANGELOG.rst

@@ -47,6 +47,7 @@ Features:
 - ``ContainerFormat.fixed_framesize`` reports whether a format wants fixed size audio frames.
 - :class:`.CodecContext` exposes more of ``AVCodecContext``: ``pkt_timebase``, ``frame_num``, ``active_thread_type``, ``bits_per_raw_sample``, ``compression_level``, ``rc_buffer_size``, ``min_bit_rate``, a setter for ``max_bit_rate``, the audio ``initial_padding``, ``trailing_padding``, and ``seek_preroll``, and ``stats_in``/``stats_out`` for two-pass encoding. ``VideoCodecContext`` gains ``chroma_sample_location``, ``refs``, and ``mb_decision``; ``AudioCodecContext`` gains ``block_align``.
 - ``CodecContext.coded_side_data`` and ``CodecContext.decoded_side_data`` expose the context's global side data as dicts of ``bytes``, keyed by packet side data name and :class:`~av.sidedata.sidedata.Type` respectively. Stream wide HDR metadata, such as mastering display and content light level, arrives in ``decoded_side_data`` once a frame has been decoded.
+- ``VideoFrame.chroma_location`` exposes ``AVFrame.chroma_location``, the position of the chroma samples relative to the luma samples, and the new ``ChromaLocation`` enum names its values. Only the codec context side of the field was wrapped, as ``VideoCodecContext.chroma_sample_location``, so the siting a decoder actually reported per frame could not be read at all. Each property mirrors its C field name, which FFmpeg spells differently on the two structs.
 - Enums gained the members FFmpeg has since added: ``Properties.FIELDS``, ``Properties.ENHANCEMENT``, ``PixFmtLoss.EXCESS_RESOLUTION``, ``PixFmtLoss.EXCESS_DEPTH``, ``Flags2.icc_profiles``, ``format.Flags.experimental``, ``Interpolation.STRICT``, ``Interpolation.UNSTABLE``, ``ColorTrc.V_LOG``, ``ColorPrimaries.V_GAMUT``, the ``LCEVC``, ``VIEW_ID``, ``THREE_D_REFERENCE_DISPLAYS``, and ``EXIF`` members of ``sidedata.Type``, and the ``exif``, ``dynamic_hdr_smpte_2094_app5``, and ``hevc_conf`` packet side data names.
 
 Fixes:

av/frame.py

@@ -165,7 +165,8 @@ def is_corrupt(self):
     def key_frame(self):
         """Is this frame a key frame?
 
-        Wraps :ffmpeg:`AVFrame.key_frame`.
+        Reads the ``AV_FRAME_FLAG_KEY`` bit of :ffmpeg:`AVFrame.flags`. FFmpeg
+        removed the ``AVFrame.key_frame`` field this used to wrap.
 
         """
         return bool(self.ptr.flags & lib.AV_FRAME_FLAG_KEY)

av/video/codeccontext.py

@@ -384,9 +384,10 @@ def coded_height(self):
     @property
     def color_range(self):
         """
-        Describes the signal range of the colorspace.
+        Describes the signal range of the colorspace, as FFmpeg's raw integer
+        value. :class:`.ColorRange` names the values.
 
-        Wraps :ffmpeg:`AVFrame.color_range`.
+        Wraps :ffmpeg:`AVCodecContext.color_range`.
 
         :type: int
         """
@@ -399,9 +400,10 @@ def color_range(self, value):
     @property
     def color_primaries(self):
         """
-        Describes the RGB/XYZ matrix of the colorspace.
+        Describes the RGB/XYZ matrix of the colorspace, as FFmpeg's raw integer
+        value. :class:`.ColorPrimaries` names the values.
 
-        Wraps :ffmpeg:`AVFrame.color_primaries`.
+        Wraps :ffmpeg:`AVCodecContext.color_primaries`.
 
         :type: int
         """
@@ -414,9 +416,11 @@ def color_primaries(self, value):
     @property
     def color_trc(self):
         """
-        Describes the linearization function (a.k.a. transformation characteristics) of the colorspace.
+        Describes the linearization function (a.k.a. transformation
+        characteristics) of the colorspace, as FFmpeg's raw integer value.
+        :class:`.ColorTrc` names the values.
 
-        Wraps :ffmpeg:`AVFrame.color_trc`.
+        Wraps :ffmpeg:`AVCodecContext.color_trc`.
 
         :type: int
         """
@@ -429,9 +433,10 @@ def color_trc(self, value):
     @property
     def colorspace(self):
         """
-        Describes the YUV/RGB transformation matrix of the colorspace.
+        Describes the YUV/RGB transformation matrix of the colorspace, as
+        FFmpeg's raw integer value. :class:`.Colorspace` names the values.
 
-        Wraps :ffmpeg:`AVFrame.colorspace`.
+        Wraps :ffmpeg:`AVCodecContext.colorspace`.
 
         :type: int
         """
@@ -500,10 +505,12 @@ def qmax(self, value):
     @property
     def chroma_sample_location(self):
         """
-        Location of the chroma samples relative to the luma samples, as
-        FFmpeg's raw integer value.
+        The position of the chroma samples relative to the luma samples, as
+        FFmpeg's raw integer value. :class:`.ChromaLocation` names the values.
 
-        Wraps :ffmpeg:`AVCodecContext.chroma_sample_location`.
+        Wraps :ffmpeg:`AVCodecContext.chroma_sample_location`. FFmpeg spells
+        the same field ``chroma_location`` on a frame, and so does PyAV: see
+        :attr:`.VideoFrame.chroma_location`.
 
         :type: int
         """

av/video/frame.py

@@ -705,10 +705,12 @@ def pict_type(self, value):
 
     @property
     def colorspace(self):
-        """Colorspace of frame.
+        """The YUV/RGB transformation matrix of the frame, as FFmpeg's raw
+        integer value. :class:`.Colorspace` names the values.
 
         Wraps :ffmpeg:`AVFrame.colorspace`.
 
+        :type: int
         """
         return self.ptr.colorspace
 
@@ -718,10 +720,12 @@ def colorspace(self, value):
 
     @property
     def color_range(self):
-        """Color range of frame.
+        """The signal range of the frame, as FFmpeg's raw integer value.
+        :class:`.ColorRange` names the values.
 
         Wraps :ffmpeg:`AVFrame.color_range`.
 
+        :type: int
         """
         return self.ptr.color_range
 
@@ -731,10 +735,13 @@ def color_range(self, value):
 
     @property
     def color_trc(self):
-        """Transfer characteristic of frame.
+        """The linearization function (a.k.a. transfer characteristic) of the
+        frame, as FFmpeg's raw integer value. :class:`.ColorTrc` names the
+        values.
 
         Wraps :ffmpeg:`AVFrame.color_trc`.
 
+        :type: int
         """
         return self.ptr.color_trc
 
@@ -744,17 +751,36 @@ def color_trc(self, value):
 
     @property
     def color_primaries(self):
-        """Color primaries of frame.
+        """The RGB/XYZ matrix of the frame, as FFmpeg's raw integer value.
+        :class:`.ColorPrimaries` names the values.
 
         Wraps :ffmpeg:`AVFrame.color_primaries`.
 
+        :type: int
         """
         return self.ptr.color_primaries
 
     @color_primaries.setter
     def color_primaries(self, value):
         self.ptr.color_primaries = value
 
+    @property
+    def chroma_location(self):
+        """The position of the chroma samples relative to the luma samples, as
+        FFmpeg's raw integer value. :class:`.ChromaLocation` names the values.
+
+        Wraps :ffmpeg:`AVFrame.chroma_location`. FFmpeg spells the same field
+        ``chroma_sample_location`` on a codec context, and so does PyAV: see
+        :attr:`.VideoCodecContext.chroma_sample_location`.
+
+        :type: int
+        """
+        return self.ptr.chroma_location
+
+    @chroma_location.setter
+    def chroma_location(self, value):
+        self.ptr.chroma_location = value
+
     def reformat(self, *args, **kwargs):
         """reformat(width=None, height=None, format=None, src_colorspace=None, dst_colorspace=None, interpolation=None, threads=None)
 

av/video/frame.pyi

@@ -54,6 +54,7 @@ class VideoFrame(Frame):
     color_range: int
     color_trc: int
     color_primaries: int
+    chroma_location: int
 
     @property
     def sw_format(self) -> VideoFormat | None: ...

av/video/reformatter.py

@@ -104,6 +104,21 @@ class ColorPrimaries(IntEnum):
     V_GAMUT: "Panasonic V-Gamut (not part of H.273)" = lib.AVCOL_PRI_V_GAMUT
 
 
+class ChromaLocation(IntEnum):
+    """Location of the chroma samples relative to the luma samples.
+
+    Maps to FFmpeg's ``AVChromaLocation``.
+    """
+
+    UNSPECIFIED: "Unspecified" = lib.AVCHROMA_LOC_UNSPECIFIED
+    LEFT: "MPEG-2/4 4:2:0, H.264 default for 4:2:0" = lib.AVCHROMA_LOC_LEFT
+    CENTER: "MPEG-1 4:2:0, JPEG 4:2:0, H.263 4:2:0" = lib.AVCHROMA_LOC_CENTER
+    TOPLEFT: "ITU-R 601, SMPTE 274M/296M, MPEG-2 4:2:2" = lib.AVCHROMA_LOC_TOPLEFT
+    TOP: "Top" = lib.AVCHROMA_LOC_TOP
+    BOTTOMLEFT: "Bottom left" = lib.AVCHROMA_LOC_BOTTOMLEFT
+    BOTTOM: "Bottom" = lib.AVCHROMA_LOC_BOTTOM
+
+
 @cython.cfunc
 @cython.inline
 def _resolve_enum_value(

av/video/reformatter.pyi

@@ -84,6 +84,15 @@ class ColorPrimaries(IntEnum):
     EBU3213 = cast(int, ...)
     V_GAMUT = cast(int, ...)
 
+class ChromaLocation(IntEnum):
+    UNSPECIFIED = cast(int, ...)
+    LEFT = cast(int, ...)
+    CENTER = cast(int, ...)
+    TOPLEFT = cast(int, ...)
+    TOP = cast(int, ...)
+    BOTTOMLEFT = cast(int, ...)
+    BOTTOM = cast(int, ...)
+
 class VideoReformatter:
     def reformat(
         self,

av/video/stream.pyi

@@ -52,4 +52,5 @@ class VideoStream(Stream):
     color_primaries: int
     color_trc: int
     colorspace: int
+    chroma_sample_location: int
     type: Literal["video"]

docs/api/video.rst

@@ -72,6 +72,28 @@ Types
     .. enumtable:: av.video.frame.PictureType
 
 
+Colors
+~~~~~~
+
+These describe how to interpret the frame's samples. They are FFmpeg's raw
+integer values, and each is named by an enum under :ref:`video_enums`. A
+decoder fills them in from the stream; an encoder passes them through to the
+container. Setting one relabels the frame, it does not convert the pixels --
+see :meth:`VideoFrame.reformat` for that.
+
+.. autoattribute:: VideoFrame.colorspace
+.. autoattribute:: VideoFrame.color_range
+.. autoattribute:: VideoFrame.color_trc
+.. autoattribute:: VideoFrame.color_primaries
+.. autoattribute:: VideoFrame.chroma_location
+
+The matching :class:`.VideoCodecContext` attributes carry the same values for a
+whole stream. FFmpeg spells the chroma one ``chroma_location`` on a frame and
+``chroma_sample_location`` on a codec context, and PyAV mirrors each C field
+name, so :attr:`.VideoCodecContext.chroma_sample_location` is the codec context
+spelling of :attr:`VideoFrame.chroma_location`.
+
+
 Conversions
 ~~~~~~~~~~~
 
@@ -105,6 +127,8 @@ Video Reformatters
 
         .. automethod:: reformat
 
+.. _video_enums:
+
 Enums
 ~~~~~
 
@@ -123,7 +147,24 @@ Enums
 
 .. autoclass:: av.video.reformatter.ColorRange
 
-    Wraps the ``AVCOL*`` flags.
+    Wraps the ``AVCOL_RANGE_*`` flags.
 
     .. enumtable:: av.video.reformatter.ColorRange
 
+.. autoclass:: av.video.reformatter.ColorTrc
+
+    Wraps the ``AVCOL_TRC_*`` flags.
+
+    .. enumtable:: av.video.reformatter.ColorTrc
+
+.. autoclass:: av.video.reformatter.ColorPrimaries
+
+    Wraps the ``AVCOL_PRI_*`` flags.
+
+    .. enumtable:: av.video.reformatter.ColorPrimaries
+
+.. autoclass:: av.video.reformatter.ChromaLocation
+
+    Wraps the ``AVCHROMA_LOC_*`` flags.
+
+    .. enumtable:: av.video.reformatter.ChromaLocation

docs/conf.py

@@ -265,6 +265,27 @@ def ffmpeg_role(name, rawtext, text, lineno, inliner, options={}, content=[]):
             "AVFrame.color_primaries": "#a59a3f830494f2ed1133103a1bc9481e7",
             "AVFrame.color_trc": "#ab09abb126e3922bc1d010cf044087939",
             "AVFrame.colorspace": "#a9262c231f1f64869439b4fe587fe1710",
+            "AVFrame.chroma_location": "#a1d15617172d8123a66bdcf8d4d826ee2",
+            "AVCodecContext.color_range": "#a255bf7100a4ba6dcb6ee5d87740a4f35",
+            "AVCodecContext.color_primaries": "#a3a41b3e5bde23b877799f6e72dac8ef3",
+            "AVCodecContext.color_trc": "#ab649e8c599f5a0e2a30448e67a36deb6",
+            "AVCodecContext.colorspace": "#a8cd8caa7d40319324ce3d879a2edbd9f",
+            "AVCodecContext.chroma_sample_location": "#ac60a0209642b5d74068cab0ac35a78b2",
+            "AVCodecContext.field_order": "#a5d222eeeb0b54ab462af363bcb9273bc",
+            "AVCodecContext.block_align": "#ae56433cc80666ff63af59db4de5b5e45",
+            "AVCodecContext.mb_decision": "#a66af0e26734255f1eacabd7d67558482",
+            "AVCodecContext.rc_max_rate": "#aa2b5582f1a360534310b686cc3f7c668",
+            "AVCodecContext.refs": "#aa0cb7241b4624dba761c8cf58fb2d5f0",
+            "AVCodecContext.time_base": "#ab7bfeb9fa5840aac090e2b0bd0ef7589",
+            "AVFrame.pict_type": "#af9920fc3fbfa347b8943ae461b50d18b",
+            "AVFrame.flags": "#a49020cc320b8fb1f5449167b6c97515b",
+            "AVFrame.pts": "#a0452833e3ab6ddd7acbf82817a7818a4",
+            "AVPacket.pts": "#a73bde0a37f3b1efc839f11295bfbf42a",
+            "AVStream.r_frame_rate": "#ad63fb11cc1415e278e09ddc676e8a1ad",
+            "AVStream.time_base": "#a9db755451f14e2bf590d4b85d82b32e6",
+            "AVFormatContext.flags": "#a32379cc371463b235d54235d4af06a15",
+            "AVFormatContext.start_time_realtime": "#aa5ddb5cee1df28f21739133f2e37f1c5",
+            "AVFilterGraph.nb_threads": "#ac28dcbf76e6fdd800295a2738d41660e",
         }.get(text, f"#{member}")
 
         url = base_url.format(struct_name) + fragment

include/avcodec.pxd

@@ -383,6 +383,7 @@ cdef extern from "libavcodec/avcodec.h" nogil:
         AVColorPrimaries color_primaries
         AVColorTransferCharacteristic color_trc
         AVColorSpace colorspace
+        AVChromaLocation chroma_location
 
         AVDictionary *metadata
         int decode_error_flags

tests/test_colorspace.py

@@ -1,7 +1,10 @@
+import io
+
 import pytest
 
 import av
 from av.video.reformatter import (
+    ChromaLocation,
     ColorPrimaries,
     ColorRange,
     Colorspace,
@@ -67,6 +70,55 @@ def test_frame_color_primaries_property() -> None:
     assert frame.color_primaries == 1  # AVCOL_PRI_BT709
 
 
+def test_frame_chroma_location_property() -> None:
+    frame = av.VideoFrame(width=64, height=64, format="yuv420p")
+    assert frame.chroma_location == ChromaLocation.UNSPECIFIED
+
+    frame.chroma_location = ChromaLocation.LEFT
+    assert frame.chroma_location == ChromaLocation.LEFT
+    assert frame.chroma_location == 1  # AVCHROMA_LOC_LEFT
+
+    frame.chroma_location = ChromaLocation.TOPLEFT
+    assert frame.chroma_location == ChromaLocation.TOPLEFT
+
+
+def test_codec_context_chroma_sample_location_property() -> None:
+    ctx = av.codec.CodecContext.create("ffv1", "w")
+    assert isinstance(ctx, av.video.codeccontext.VideoCodecContext)
+    assert ctx.chroma_sample_location == ChromaLocation.UNSPECIFIED
+
+    ctx.chroma_sample_location = ChromaLocation.CENTER
+    assert ctx.chroma_sample_location == ChromaLocation.CENTER
+    assert ctx.chroma_sample_location == 2  # AVCHROMA_LOC_CENTER
+
+
+def test_chroma_location_round_trip() -> None:
+    # ffv1 in matroska signals the chroma location verbatim.
+    buf = io.BytesIO()
+    with av.open(buf, "w", format="matroska") as output:
+        out_stream = output.add_stream("ffv1", rate=30)
+        out_stream.width = 64
+        out_stream.height = 64
+        out_stream.pix_fmt = "yuv420p"
+        out_stream.codec_context.chroma_sample_location = ChromaLocation.TOPLEFT
+
+        for i in range(5):
+            frame = av.VideoFrame(width=64, height=64, format="yuv420p")
+            frame.pts = i
+            output.mux(out_stream.encode(frame))
+        output.mux(out_stream.encode())
+
+    buf.seek(0)
+    with av.open(buf, "r") as input_:
+        in_stream = input_.streams.video[0]
+        assert in_stream.chroma_sample_location == ChromaLocation.TOPLEFT
+
+        frames = list(input_.decode(in_stream))
+        assert frames
+        for decoded in frames:
+            assert decoded.chroma_location == ChromaLocation.TOPLEFT
+
+
 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")