From 003623bae5a43c2d0e2a1c893ac3e8f2f795d6a8 Mon Sep 17 00:00:00 2001 From: quietvoid <39477805+quietvoid@users.noreply.github.com> Date: Thu, 25 Aug 2022 08:41:42 -0400 Subject: [PATCH] Implement support for encoding HDR10+ from JSON metadata file --- Cargo.lock | 23 ++++++++++++++ Cargo.toml | 1 + src/api/config/encoder.rs | 7 +++++ src/api/internal.rs | 65 +++++++++++++++++++++++++++++++++++++-- src/api/test.rs | 2 ++ src/api/util.rs | 8 +++++ src/bin/common.rs | 37 ++++++++++++++++++++++ src/fuzzing.rs | 1 + 8 files changed, 142 insertions(+), 2 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 2939c9ea9b..fac3256e97 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -185,6 +185,15 @@ version = "1.6.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9d28070975aaf4ef1fd0bd1f29b739c06c2cdd9972e090617fb6dca3b2cb564e" +[[package]] +name = "bitvec_helpers" +version = "3.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3ef6883bd86b4112b56be19de3a1628de6c4063be7be6e641d484c83069efb4a" +dependencies = [ + "bitstream-io", +] + [[package]] name = "bstr" version = "1.1.0" @@ -684,6 +693,18 @@ version = "0.12.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8a9ee70c43aaf417c914396645a0fa852624801b24ebb7ae78fe8272889ac888" +[[package]] +name = "hdr10plus" +version = "2.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ee3a6bfa9ada73834523412702bbcaa0f53c3393acb36278cd3489c38c808491" +dependencies = [ + "anyhow", + "bitvec_helpers", + "serde", + "serde_json", +] + [[package]] name = "heck" version = "0.4.0" @@ -1305,6 +1326,7 @@ dependencies = [ "crossbeam", "dav1d-sys", "fern", + "hdr10plus", "image", "interpolate_name", "itertools 0.10.5", @@ -1531,6 +1553,7 @@ version = "1.0.91" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "877c235533714907a8c2464236f5c4b2a17262ef1bd71f38f35ea592c8da6883" dependencies = [ + "indexmap", "itoa", "ryu", "serde", diff --git a/Cargo.toml b/Cargo.toml index 4d14f10a60..5d63c3fb54 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -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" diff --git a/src/api/config/encoder.rs b/src/api/config/encoder.rs index 038f3301b0..15454b255a 100644 --- a/src/api/config/encoder.rs +++ b/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. @@ -91,6 +92,11 @@ pub struct EncoderConfig { pub tune: Tune, /// Parameters for grain synthesis. pub film_grain_params: Option>, + /// 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>>, /// 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, diff --git a/src/api/internal.rs b/src/api/internal.rs index 20ef57e328..3aee5f642b 100644 --- a/src/api/internal.rs +++ b/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,6 +350,12 @@ impl ContextInner { } 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); @@ -356,7 +363,14 @@ impl ContextInner { 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) { @@ -1725,4 +1739,51 @@ impl ContextInner { (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> { + 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 + } } diff --git a/src/api/test.rs b/src/api/test.rs index 5d08692ffe..7e417607ec 100644 --- a/src/api/test.rs +++ b/src/api/test.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, diff --git a/src/api/util.rs b/src/api/util.rs index a2ab9794e6..19475c2a7c 100644 --- a/src/api/util.rs +++ b/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 { diff --git a/src/bin/common.rs b/src/bin/common.rs index 2c9b120365..653791a506 100644 --- a/src/bin/common.rs +++ b/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::*; @@ -194,6 +195,14 @@ pub struct CliOptions { help_heading = "ENCODE SETTINGS" )] pub film_grain_table: Option, + /// 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, /// Pixel range #[clap(long, value_parser, help_heading = "VIDEO METADATA")] @@ -678,6 +687,34 @@ fn parse_config(matches: &CliOptions) -> Result { } } + 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> = 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); } diff --git a/src/fuzzing.rs b/src/fuzzing.rs index d6a7b5e6ad..790b8dc6cf 100644 --- a/src/fuzzing.rs +++ b/src/fuzzing.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 =