Implement support for encoding Dolby Vision from RPU file

Author: quietvoid

Cargo.lock

@@ -108,6 +108,7 @@ new_debug_unreachable = "1.0.4"
 once_cell = "1.18.0"
 av1-grain = { version = "0.2.2", features = ["serialize"] }
 serde-big-array = { version = "0.5.1", optional = true }
+dolby_vision = { version = "3.2.0" }
 
 [dependencies.image]
 version = "0.24.6"

Cargo.toml

@@ -9,12 +9,13 @@
 
 use itertools::*;
 
-use crate::api::color::*;
 use crate::api::config::GrainTableSegment;
+use crate::api::{color::*, T35};
 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,8 @@ pub struct EncoderConfig {
   pub tune: Tune,
   /// Parameters for grain synthesis.
   pub film_grain_params: Option<Vec<GrainTableSegment>>,
+  /// Dolby Vision T.35 metadata payload map, by input frame index.
+  pub dovi_payloads: Option<BTreeMap<u64, T35>>,
   /// Number of tiles horizontally. Must be a power of two.
   ///
   /// Overridden by [`tiles`], if present.
@@ -167,6 +170,7 @@ impl EncoderConfig {
       bitrate: 0,
       tune: Tune::default(),
       film_grain_params: None,
+      dovi_payloads: None,
       tile_cols: 0,
       tile_rows: 0,
       tiles: 0,

src/api/config/encoder.rs

@@ -2127,6 +2127,7 @@ fn log_q_exp_overflow() {
     bitrate: 1,
     tune: Tune::Psychovisual,
     film_grain_params: None,
+    dovi_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,
+    dovi_payloads: None,
     tile_cols: 0,
     tile_rows: 0,
     tiles: 0,

src/api/test.rs

@@ -137,6 +137,12 @@ impl fmt::Display for FrameType {
   }
 }
 
+/// Dolby Vision T.35 metadata payload expected prefix.
+pub const T35_DOVI_PAYLOAD_PREFIX: &[u8] = &[
+  0x00, 0x03B, // Dolby
+  0x00, 0x00, 0x08, 0x00, 0x37, 0xCD, 0x08,
+];
+
 /// A single T.35 metadata packet.
 #[derive(Clone, Debug, Default)]
 pub struct T35 {
@@ -299,3 +305,10 @@ 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 Dolby Vision Metadata.
+  pub fn is_dovi_metadata(&self) -> bool {
+    self.country_code == 0xB5 && self.data.starts_with(T35_DOVI_PAYLOAD_PREFIX)
+  }
+}

src/api/util.rs

@@ -18,6 +18,8 @@ use rav1e::prelude::*;
 use scan_fmt::scan_fmt;
 
 use rav1e::config::CpuFeatureLevel;
+
+use std::collections::BTreeMap;
 use std::fs::File;
 use std::io;
 use std::io::prelude::*;
@@ -195,6 +197,15 @@ pub struct CliOptions {
     help_heading = "ENCODE SETTINGS"
   )]
   pub film_grain_table: Option<PathBuf>,
+  /// Uses a Dolby Vision RPU file to add as T.35 metadata to the encode.
+  /// The RPU must be in the same format as for x265
+  #[clap(
+    long,
+    alias = "dolby-vision-rpu",
+    value_parser,
+    help_heading = "ENCODE SETTINGS"
+  )]
+  pub dovi_rpu: Option<PathBuf>,
 
   /// Pixel range
   #[clap(long, value_parser, help_heading = "VIDEO METADATA")]
@@ -684,6 +695,31 @@ fn parse_config(matches: &CliOptions) -> Result<EncoderConfig, CliError> {
     }
   }
 
+  if let Some(rpu_file) = matches.dovi_rpu.as_ref() {
+    let rpus = dolby_vision::rpu::utils::parse_rpu_file(rpu_file)
+      .expect("Failed to read Dolby Vision RPU file");
+
+    let payloads: BTreeMap<u64, T35> = rpus
+      .iter()
+      .filter_map(|rpu| {
+        rpu
+          .write_av1_rpu_metadata_obu_t35_payload()
+          .map(|payload| T35 {
+            country_code: 0xB5,
+            country_code_extension_byte: 0x00,
+            data: payload.into_boxed_slice(),
+          })
+          .ok()
+      })
+      .zip(0u64..)
+      .map(|(payload, frame_no)| (frame_no, payload))
+      .collect();
+
+    if !payloads.is_empty() {
+      cfg.dovi_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

@@ -1284,6 +1284,19 @@ impl<T: Pixel> FrameInvariants<T> {
     self.input_frameno * TIMESTAMP_BASE_UNIT * self.sequence.time_base.num
       / self.sequence.time_base.den
   }
+
+  /// Dolby Vision Metadata as T.35 metadata from [`EncoderConfig`]
+  pub fn dovi_metadata(&self) -> Option<&T35> {
+    if !(self.show_frame || self.is_show_existing_frame()) {
+      return None;
+    }
+
+    self
+      .config
+      .dovi_payloads
+      .as_ref()
+      .and_then(|payloads| payloads.get(&self.input_frameno))
+  }
 }
 
 impl<T: Pixel> fmt::Display for FrameInvariants<T> {
@@ -3686,11 +3699,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.dovi_metadata() {
+    if !fi.t35_metadata.iter().any(|t35| t35.is_dovi_metadata()) {
+      write_t35_metadata_packet(&mut packet, t35);
+    }
   }
 
   let mut buf1 = Vec::new();
@@ -3767,11 +3783,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.dovi_metadata() {
+    if !fi.t35_metadata.iter().any(|t35| t35.is_dovi_metadata()) {
+      write_t35_metadata_packet(&mut packet, t35);
+    }
   }
 
   let mut buf1 = Vec::new();
@@ -3827,6 +3846,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,
+      dovi_payloads: None,
     };
 
     let frame_count =