Expose HDR Color metadata (#1271)

Author: mollyxu

src/torchcodec/_core/Metadata.cpp

@@ -7,6 +7,10 @@
 #include "Metadata.h"
 #include "StableABICompat.h"
 
+extern "C" {
+#include <libavutil/pixdesc.h>
+}
+
 namespace facebook::torchcodec {
 
 std::optional<double> StreamMetadata::getDurationSeconds(
@@ -123,4 +127,38 @@ std::optional<double> StreamMetadata::getAverageFps(SeekMode seekMode) const {
   }
 }
 
+std::optional<std::string> StreamMetadata::getColorPrimariesName() const {
+  if (!colorPrimaries.has_value()) {
+    return std::nullopt;
+  }
+  const char* name = av_color_primaries_name(*colorPrimaries);
+  if (name == nullptr) {
+    return std::nullopt;
+  }
+  return std::string(name);
+}
+
+std::optional<std::string> StreamMetadata::getColorSpaceName() const {
+  if (!colorSpace.has_value()) {
+    return std::nullopt;
+  }
+  const char* name = av_color_space_name(*colorSpace);
+  if (name == nullptr) {
+    return std::nullopt;
+  }
+  return std::string(name);
+}
+
+std::optional<std::string> StreamMetadata::getColorTransferCharacteristicName()
+    const {
+  if (!colorTransferCharacteristic.has_value()) {
+    return std::nullopt;
+  }
+  const char* name = av_color_transfer_name(*colorTransferCharacteristic);
+  if (name == nullptr) {
+    return std::nullopt;
+  }
+  return std::string(name);
+}
+
 } // namespace facebook::torchcodec

src/torchcodec/_core/Metadata.h

@@ -13,6 +13,7 @@
 extern "C" {
 #include <libavcodec/avcodec.h>
 #include <libavutil/avutil.h>
+#include <libavutil/pixfmt.h>
 #include <libavutil/rational.h>
 }
 
@@ -59,6 +60,9 @@ struct StreamMetadata {
   std::optional<AVRational> sampleAspectRatio;
   // Rotation angle in degrees from display matrix, in the range [-180, 180].
   std::optional<double> rotation;
+  std::optional<AVColorPrimaries> colorPrimaries;
+  std::optional<AVColorSpace> colorSpace;
+  std::optional<AVColorTransferCharacteristic> colorTransferCharacteristic;
   // The pixel format of the encoded video, e.g. "yuv420p".
   std::optional<std::string> pixelFormat;
 
@@ -73,6 +77,12 @@ struct StreamMetadata {
   std::optional<double> getEndStreamSeconds(SeekMode seekMode) const;
   std::optional<int64_t> getNumFrames(SeekMode seekMode) const;
   std::optional<double> getAverageFps(SeekMode seekMode) const;
+
+  // Color metadata name accessors. These return nullopt if the field is unset
+  // or if FFmpeg returns NULL for the name.
+  std::optional<std::string> getColorPrimariesName() const;
+  std::optional<std::string> getColorSpaceName() const;
+  std::optional<std::string> getColorTransferCharacteristicName() const;
 };
 
 struct ContainerMetadata {

src/torchcodec/_core/SingleStreamDecoder.cpp

@@ -171,6 +171,16 @@ void SingleStreamDecoder::initializeDecoder() {
       streamMetadata.sampleAspectRatio =
           avStream->codecpar->sample_aspect_ratio;
 
+      if (avStream->codecpar->color_primaries != AVCOL_PRI_UNSPECIFIED) {
+        streamMetadata.colorPrimaries = avStream->codecpar->color_primaries;
+      }
+      if (avStream->codecpar->color_space != AVCOL_SPC_UNSPECIFIED) {
+        streamMetadata.colorSpace = avStream->codecpar->color_space;
+      }
+      if (avStream->codecpar->color_trc != AVCOL_TRC_UNSPECIFIED) {
+        streamMetadata.colorTransferCharacteristic =
+            avStream->codecpar->color_trc;
+      }
       AVPixelFormat pixelFormat =
           static_cast<AVPixelFormat>(avStream->codecpar->format);
       // If the AVPixelFormat is not recognized, we get back nullptr. We have

src/torchcodec/_core/_metadata.py

@@ -117,8 +117,17 @@ class VideoStreamMetadata(StreamMetadata):
         or -90 degree rotation, this means width and height are swapped
         compared to the raw encoded dimensions in the container.
     """
+    color_primaries: str | None
+    """Color primaries as reported by FFmpeg. E.g. ``"bt709"``, ``"bt2020"``."""
+    color_space: str | None
+    """Color space as reported by FFmpeg. E.g. ``"bt709"``,
+    ``"bt2020nc"``."""
+    color_transfer_characteristic: str | None
+    """Color transfer characteristic as reported by FFmpeg
+    E.g. ``"bt709"``, ``"smpte2084"`` (PQ), ``"arib-std-b67"`` (HLG)."""
     pixel_format: str | None
-    """The source pixel format of the video, as reported by FFmpeg. E.g. ``'yuv420p'``, ``'yuv444p'``, etc."""
+    """The source pixel format of the video as reported by FFmpeg.
+    E.g. ``'yuv420p'``, ``'yuv444p'``, etc."""
 
     # Computed fields (computed in C++ with fallback logic)
     end_stream_seconds: float | None
@@ -245,6 +254,11 @@ def get_container_metadata(decoder: torch.Tensor) -> ContainerMetadata:
                     average_fps_from_header=stream_dict.get("averageFpsFromHeader"),
                     pixel_aspect_ratio=_get_optional_par_fraction(stream_dict),
                     rotation=stream_dict.get("rotation"),
+                    color_primaries=stream_dict.get("colorPrimaries"),
+                    color_space=stream_dict.get("colorSpace"),
+                    color_transfer_characteristic=stream_dict.get(
+                        "colorTransferCharacteristic"
+                    ),
                     pixel_format=stream_dict.get("pixelFormat"),
                     **common_meta,
                 )

src/torchcodec/_core/custom_ops.cpp

@@ -9,6 +9,11 @@
 #include <cstdint>
 #include <sstream>
 #include <string>
+
+extern "C" {
+#include <libavutil/pixdesc.h>
+}
+
 #include "AVIOFileLikeContext.h"
 #include "AVIOTensorContext.h"
 #include "Encoder.h"
@@ -987,6 +992,15 @@ std::string get_stream_json_metadata(
   if (streamMetadata.rotation.has_value()) {
     map["rotation"] = std::to_string(*streamMetadata.rotation);
   }
+  if (auto name = streamMetadata.getColorPrimariesName()) {
+    map["colorPrimaries"] = quoteValue(*name);
+  }
+  if (auto name = streamMetadata.getColorSpaceName()) {
+    map["colorSpace"] = quoteValue(*name);
+  }
+  if (auto name = streamMetadata.getColorTransferCharacteristicName()) {
+    map["colorTransferCharacteristic"] = quoteValue(*name);
+  }
   if (streamMetadata.pixelFormat.has_value()) {
     map["pixelFormat"] = quoteValue(streamMetadata.pixelFormat.value());
   }

test/test_metadata.py

@@ -8,7 +8,6 @@
 from fractions import Fraction
 
 import pytest
-
 from torchcodec import ffmpeg_major_version
 from torchcodec._core import (
     add_video_stream,
@@ -20,7 +19,12 @@
 )
 from torchcodec.decoders import AudioDecoder, VideoDecoder
 
-from .utils import NASA_AUDIO_MP3, NASA_VIDEO, NASA_VIDEO_ROTATED
+from .utils import (
+    BT2020_LIMITED_RANGE_10BIT,
+    NASA_AUDIO_MP3,
+    NASA_VIDEO,
+    NASA_VIDEO_ROTATED,
+)
 
 
 # TODO: Expected values in these tests should be based on the assets's
@@ -169,6 +173,22 @@ def test_rotation_metadata():
     )
 
 
+def test_color_metadata():
+    # BT2020_LIMITED_RANGE_10BIT is a BT.2020 10-bit HEVC video with PQ transfer
+    decoder_bt2020 = VideoDecoder(BT2020_LIMITED_RANGE_10BIT.path)
+    assert decoder_bt2020.metadata.color_primaries == "bt2020"
+    assert decoder_bt2020.metadata.color_space == "bt2020nc"
+    assert decoder_bt2020.metadata.color_transfer_characteristic == "smpte2084"
+    assert decoder_bt2020.metadata.pixel_format == "yuv420p10le"
+
+    # NASA_VIDEO has BT.709 color metadata
+    decoder_nasa = VideoDecoder(NASA_VIDEO.path)
+    assert decoder_nasa.metadata.color_primaries == "bt709"
+    assert decoder_nasa.metadata.color_space == "bt709"
+    assert decoder_nasa.metadata.color_transfer_characteristic == "bt709"
+    assert decoder_nasa.metadata.pixel_format == "yuv420p"
+
+
 def test_repr():
     # Test for calls to print(), str(), etc. Useful to make sure we don't forget
     # to add additional @properties to __repr__
@@ -191,6 +211,9 @@ def test_repr():
   average_fps_from_header: 29.97002997002997
   pixel_aspect_ratio: 1
   rotation: None
+  color_primaries: bt709
+  color_space: bt709
+  color_transfer_characteristic: bt709
   pixel_format: yuv420p
   end_stream_seconds: 13.013
   num_frames: 390