Implement support for encoding HDR10+ from JSON metadata file

Author: quietvoid

Cargo.toml

@@ -108,6 +108,7 @@ new_debug_unreachable = "1.0.4"
 once_cell = "1.13.0"
 av1-grain = { version = "0.2.0", features = ["serialize"] }
 serde-big-array = { version = "0.4.1", optional = true }
+hdr10plus = { version = "1.2.0", features = ["json"] }
 
 [dependencies.image]
 version = "0.24.3"

src/api/config/encoder.rs

@@ -15,6 +15,7 @@ use crate::api::{Rational, SpeedSettings};
 use crate::encoder::Tune;
 use crate::serialize::{Deserialize, Serialize};
 
+use std::collections::BTreeMap;
 use std::fmt;
 
 // We add 1 to rdo_lookahead_frames in a bunch of places.
@@ -85,6 +86,11 @@ pub struct EncoderConfig {
   pub tune: Tune,
   /// Parameters for grain synthesis.
   pub film_grain_params: Option<Vec<GrainTableSegment>>,
+  /// HDR10+, ST2094-40 T.35 metadata payload map, by input frame index.
+  ///
+  /// The payloads are expected to follow the specification
+  /// defined at https://aomediacodec.github.io/av1-hdr10plus.
+  pub hdr10plus_payloads: Option<BTreeMap<u64, Vec<u8>>>,
   /// Number of tiles horizontally. Must be a power of two.
   ///
   /// Overridden by [`tiles`], if present.
@@ -159,6 +165,7 @@ impl EncoderConfig {
       bitrate: 0,
       tune: Tune::default(),
       film_grain_params: None,
+      hdr10plus_payloads: None,
       tile_cols: 0,
       tile_rows: 0,
       tiles: 0,

src/api/internal.rs

@@ -11,7 +11,8 @@
 use crate::activity::ActivityMask;
 use crate::api::lookahead::*;
 use crate::api::{
-  EncoderConfig, EncoderStatus, FrameType, Opaque, Packet, T35,
+  EncoderConfig, EncoderStatus, FrameType, Opaque, Packet, ST2094_40_PREFIX,
+  T35,
 };
 use crate::color::ChromaSampling::Cs400;
 use crate::cpu_features::CpuFeatureLevel;
@@ -349,14 +350,27 @@ impl<T: Pixel> ContextInner<T> {
     }
     self.frame_q.insert(input_frameno, frame);
 
+    // Update T.35 metadata from encoder config
+    let maybe_updated_t35_metadata = self.get_maybe_updated_t35_metadata(
+      input_frameno,
+      params.as_ref().map(|params| params.t35_metadata.as_ref()),
+    );
+
     if let Some(params) = params {
       if params.frame_type_override == FrameTypeOverride::Key {
         self.keyframes_forced.insert(input_frameno);
       }
       if let Some(op) = params.opaque {
         self.opaque_q.insert(input_frameno, op);
       }
-      self.t35_q.insert(input_frameno, params.t35_metadata);
+
+      if let Some(new_t35_metadata) = maybe_updated_t35_metadata {
+        self.t35_q.insert(input_frameno, new_t35_metadata.into_boxed_slice());
+      } else {
+        self.t35_q.insert(input_frameno, params.t35_metadata);
+      }
+    } else if let Some(new_t35_metadata) = maybe_updated_t35_metadata {
+      self.t35_q.insert(input_frameno, new_t35_metadata.into_boxed_slice());
     }
 
     if !self.needs_more_frame_q_lookahead(self.next_lookahead_frame) {
@@ -1719,4 +1733,51 @@ impl<T: Pixel> ContextInner<T> {
       (prev_keyframe_nframes, prev_keyframe_ntus)
     }
   }
+
+  /// Updates the T.35 metadata to be added to the frame.
+  /// The existing T.35 array may come from `FrameParameters`.
+  ///
+  /// Added from [`EncoderConfig`]:
+  ///   - HDR10+, ST2094-40 in `hdr10plus_payloads`.
+  ///
+  /// Returns an `Option`, where `None` means the T.35 metadata is unchanged.
+  /// Otherwise, the updated T.35 metadata is returned.
+  fn get_maybe_updated_t35_metadata(
+    &self, input_frameno: u64, maybe_existing_t35_metadata: Option<&[T35]>,
+  ) -> Option<Vec<T35>> {
+    let hdr10plus_payload = self
+      .config
+      .hdr10plus_payloads
+      .as_ref()
+      .and_then(|list| list.get(&input_frameno));
+
+    let update_t35_metadata = hdr10plus_payload.is_some();
+
+    let mut new_t35_metadata = if update_t35_metadata {
+      Some(
+        maybe_existing_t35_metadata.map_or_else(Vec::new, |t35| t35.to_vec()),
+      )
+    } else {
+      None
+    };
+
+    if let Some(list) = new_t35_metadata.as_mut() {
+      // HDR10+, ST2094-40
+      if let Some(payload) = hdr10plus_payload {
+        let has_existing_hdr10plus_meta = list.iter().any(|t35| {
+          t35.country_code == 0xB5 && t35.data.starts_with(ST2094_40_PREFIX)
+        });
+
+        if !has_existing_hdr10plus_meta {
+          list.push(T35 {
+            country_code: 0xB5,
+            country_code_extension_byte: 0x00,
+            data: payload.clone().into_boxed_slice(),
+          });
+        }
+      }
+    }
+
+    new_t35_metadata
+  }
 }

src/api/test.rs

@@ -2128,6 +2128,7 @@ fn log_q_exp_overflow() {
     bitrate: 1,
     tune: Tune::Psychovisual,
     film_grain_params: None,
+    hdr10plus_payloads: None,
     tile_cols: 0,
     tile_rows: 0,
     tiles: 0,
@@ -2204,6 +2205,7 @@ fn guess_frame_subtypes_assert() {
     bitrate: 16384,
     tune: Tune::Psychovisual,
     film_grain_params: None,
+    hdr10plus_payloads: None,
     tile_cols: 0,
     tile_rows: 0,
     tiles: 0,

src/api/util.rs

@@ -137,6 +137,14 @@ impl fmt::Display for FrameType {
   }
 }
 
+/// ST2094-40 T.35 metadata payload expected prefix.
+pub const ST2094_40_PREFIX: &[u8] = &[
+  0x00, 0x03C, // Samsung Electronics America
+  0x00, 0x01, // ST-2094-40
+  0x04, // application_identifier = 4
+  0x01, // application_mode =1
+];
+
 /// A single T.35 metadata packet.
 #[derive(Clone, Debug, Default)]
 pub struct T35 {

src/bin/common.rs

@@ -17,6 +17,7 @@ use once_cell::sync::Lazy;
 use rav1e::prelude::*;
 use scan_fmt::scan_fmt;
 
+use std::collections::BTreeMap;
 use std::fs::File;
 use std::io;
 use std::io::prelude::*;
@@ -192,6 +193,14 @@ pub struct CliOptions {
     help_heading = "ENCODE SETTINGS"
   )]
   pub film_grain_table: Option<PathBuf>,
+  /// Uses a HDR10+ metadata JSON file to add as T.35 metadata to the encode.
+  #[clap(
+    long,
+    alias = "dhdr10-info",
+    value_parser,
+    help_heading = "ENCODE SETTINGS"
+  )]
+  pub hdr10plus_json: Option<PathBuf>,
 
   /// Pixel range
   #[clap(long, value_parser, help_heading = "VIDEO METADATA")]
@@ -664,6 +673,34 @@ fn parse_config(matches: &CliOptions) -> Result<EncoderConfig, CliError> {
     }
   }
 
+  if let Some(json_file) = matches.hdr10plus_json.as_ref() {
+    let contents = std::fs::read_to_string(json_file)
+      .expect("Failed to read HDR10+ metadata file");
+    let metadata_root =
+      hdr10plus::metadata_json::MetadataJsonRoot::parse(&contents)
+        .expect("Failed to parse HDR10+ metadata");
+
+    let hdr10plus_enc_opts = hdr10plus::metadata::Hdr10PlusMetadataEncOpts {
+      with_country_code: false,
+      ..Default::default()
+    };
+    let payloads: BTreeMap<u64, Vec<u8>> = metadata_root
+      .scene_info
+      .iter()
+      .filter_map(|meta| {
+        hdr10plus::metadata::Hdr10PlusMetadata::try_from(meta)
+          .and_then(|meta| meta.encode_with_opts(&hdr10plus_enc_opts))
+          .ok()
+      })
+      .zip(0u64..)
+      .map(|(payload, frame_no)| (frame_no, payload))
+      .collect();
+
+    if !payloads.is_empty() {
+      cfg.hdr10plus_payloads = Some(payloads);
+    }
+  }
+
   if let Some(frame_rate) = matches.frame_rate {
     cfg.time_base = Rational::new(matches.time_scale, frame_rate);
   }

src/fuzzing.rs

@@ -257,6 +257,7 @@ impl Arbitrary for ArbitraryEncoder {
       switch_frame_interval: u.int_in_range(0..=3)?,
       tune: *u.choose(&[Tune::Psnr, Tune::Psychovisual])?,
       film_grain_params: None,
+      hdr10plus_payloads: None,
     };
 
     let frame_count =