Implement support for encoding HDR10+ from JSON metadata file

Author: quietvoid

Cargo.lock

@@ -109,6 +109,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 = "2.0.0", features = ["json"] }
 
 [dependencies.image]
 version = "0.24.3"

Cargo.toml

@@ -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.
@@ -91,6 +92,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.
@@ -167,6 +173,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/config/encoder.rs

@@ -2129,6 +2129,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,
@@ -2206,6 +2207,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/test.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 {
@@ -299,3 +307,11 @@ impl<T: Pixel> IntoFrame<T> for (Frame<T>, Option<FrameParameters>) {
     (Some(Arc::new(self.0)), self.1)
   }
 }
+
+impl T35 {
+  /// Whether the T.35 metadata is HDR10+ Metadata
+  /// According to the specification at https://aomediacodec.github.io/av1-hdr10plus
+  pub fn is_hdr10plus_metadata(&self) -> bool {
+    self.country_code == 0xB5 && self.data.starts_with(ST2094_40_PREFIX)
+  }
+}

src/api/util.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::*;
@@ -194,6 +195,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")]
@@ -678,6 +687,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/bin/common.rs

@@ -1286,6 +1286,26 @@ impl<T: Pixel> FrameInvariants<T> {
     self.input_frameno * TIMESTAMP_BASE_UNIT * self.sequence.time_base.num
       / self.sequence.time_base.den
   }
+
+  /// HDR10+ Metadata as T.35 metadata from [`EncoderConfig`]
+  pub fn hdr10plus_metadata(&self) -> Option<T35> {
+    if !(self.show_frame || self.is_show_existing_frame()) {
+      return None;
+    }
+
+    self
+      .config
+      .hdr10plus_payloads
+      .as_ref()
+      .and_then(|payloads| payloads.get(&self.input_frameno))
+      .and_then(|payload| {
+        Some(T35 {
+          country_code: 0xB5,
+          country_code_extension_byte: 0x00,
+          data: payload.clone().into_boxed_slice(),
+        })
+      })
+  }
 }
 
 impl<T: Pixel> fmt::Display for FrameInvariants<T> {
@@ -3681,11 +3701,14 @@ pub fn encode_show_existing_frame<T: Pixel>(
   }
 
   for t35 in fi.t35_metadata.iter() {
-    let mut t35_buf = Vec::new();
-    let mut t35_bw = BitWriter::endian(&mut t35_buf, BigEndian);
-    t35_bw.write_t35_metadata_obu(t35).unwrap();
-    packet.write_all(&t35_buf).unwrap();
-    t35_buf.clear();
+    write_t35_metadata_packet(&mut packet, t35);
+  }
+
+  // HDR10+ Metadata OBU from config
+  if let Some(t35) = fi.hdr10plus_metadata() {
+    if !fi.t35_metadata.iter().any(|t35| t35.is_hdr10plus_metadata()) {
+      write_t35_metadata_packet(&mut packet, &t35);
+    }
   }
 
   let mut buf1 = Vec::new();
@@ -3761,11 +3784,14 @@ pub fn encode_frame<T: Pixel>(
   }
 
   for t35 in fi.t35_metadata.iter() {
-    let mut t35_buf = Vec::new();
-    let mut t35_bw = BitWriter::endian(&mut t35_buf, BigEndian);
-    t35_bw.write_t35_metadata_obu(t35).unwrap();
-    packet.write_all(&t35_buf).unwrap();
-    t35_buf.clear();
+    write_t35_metadata_packet(&mut packet, t35);
+  }
+
+  // HDR10+ Metadata OBU from config
+  if let Some(t35) = fi.hdr10plus_metadata() {
+    if !fi.t35_metadata.iter().any(|t35| t35.is_hdr10plus_metadata()) {
+      write_t35_metadata_packet(&mut packet, &t35);
+    }
   }
 
   let mut buf1 = Vec::new();
@@ -3821,6 +3847,14 @@ pub fn update_rec_buffer<T: Pixel>(
   }
 }
 
+fn write_t35_metadata_packet(packet: &mut Vec<u8>, t35: &T35) {
+  let mut t35_buf = Vec::new();
+  let mut t35_bw = BitWriter::endian(&mut t35_buf, BigEndian);
+  t35_bw.write_t35_metadata_obu(t35).unwrap();
+  packet.write_all(&t35_buf).unwrap();
+  t35_buf.clear();
+}
+
 #[cfg(test)]
 mod test {
   use super::*;

src/encoder.rs

@@ -258,6 +258,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 =