Replace existing metadata when injecting

quietvoid · Jan 7, 2022 · a1d230e · a1d230e
1 parent be2c766
commit a1d230e
Show file tree

Hide file tree

Showing 3 changed files with 70 additions and 35 deletions.
diff --git a/src/commands/inject.rs b/src/commands/inject.rs
@@ -9,6 +9,8 @@ use anyhow::{bail, Result};
 const OUT_NAL_HEADER: &[u8] = &[0, 0, 0, 1];
 use hdr10plus::metadata_json::{Hdr10PlusJsonMetadata, MetadataJsonRoot};
 
+use crate::core::is_st2094_40_sei;
+
 use super::{initialize_progress_bar, input_format, Format};
 
 use hevc_parser::hevc::*;
@@ -77,6 +79,8 @@ impl Injector {
 
         let mut offsets = Vec::with_capacity(2048);
 
+        let mut already_checked_for_hdr10plus = false;
+
         while let Ok(n) = reader.read(&mut main_buf) {
             let read_bytes = n;
             if read_bytes == 0 && end.is_empty() && chunk.is_empty() {
@@ -106,7 +110,21 @@ impl Injector {
                 last
             };
 
-            parser.split_nals(&chunk, &offsets, last, true)?;
+            if !already_checked_for_hdr10plus {
+                let nals = parser.split_nals(&chunk, &offsets, last, true)?;
+
+                let contains_hdr10plus = nals.iter().any(|nal| {
+                    nal.nal_type == NAL_SEI_PREFIX
+                        && is_st2094_40_sei(&chunk[nal.start..nal.end]).unwrap_or(false)
+                });
+
+                if contains_hdr10plus {
+                    already_checked_for_hdr10plus = true;
+                    println!("\nWarning: Input file already has HDR10+ metadata SEIs, they will be replaced.");
+                }
+            } else {
+                parser.split_nals(&chunk, &offsets, last, true)?;
+            }
 
             chunk.clear();
 
@@ -118,6 +136,10 @@ impl Injector {
             consumed += read_bytes;
 
             if consumed >= 100_000_000 {
+                if !already_checked_for_hdr10plus {
+                    already_checked_for_hdr10plus = true;
+                }
+
                 pb.inc(1);
                 consumed = 0;
             }
@@ -260,8 +282,13 @@ impl Injector {
                     //    continue;
                     //}
 
-                    writer.write_all(OUT_NAL_HEADER)?;
-                    writer.write_all(&chunk[nal.start..nal.end])?;
+                    if nal.nal_type == NAL_SEI_PREFIX
+                        && is_st2094_40_sei(&chunk[nal.start..nal.end])?
+                    {
+                    } else {
+                        writer.write_all(OUT_NAL_HEADER)?;
+                        writer.write_all(&chunk[nal.start..nal.end])?;
+                    }
 
                     let global_index = nals_parsed + cur_index;
 

diff --git a/src/core/mod.rs b/src/core/mod.rs
@@ -1,6 +1,7 @@
 use anyhow::Result;
+use hevc_parser::hevc::{SeiMessage, USER_DATA_REGISTERED_ITU_T_35};
 use indicatif::{ProgressBar, ProgressStyle};
-use std::{fs::File, path::Path};
+use std::{convert::TryInto, fs::File, path::Path};
 
 use super::Format;
 
@@ -30,3 +31,36 @@ pub fn initialize_progress_bar(format: &Format, input: &Path) -> Result<Progress
 
     Ok(pb)
 }
+
+pub fn is_st2094_40_sei(sei_payload: &[u8]) -> Result<bool> {
+    if sei_payload.len() >= 4 {
+        let sei = SeiMessage::from_bytes(sei_payload)?;
+
+        if sei.payload_type == USER_DATA_REGISTERED_ITU_T_35 {
+            // FIXME: Not sure why 4 bytes..
+            let itu_t35_bytes = &sei_payload[4..];
+
+            if itu_t35_bytes.len() >= 7 {
+                let itu_t_t35_country_code = itu_t35_bytes[0];
+                let itu_t_t35_terminal_provider_code =
+                    u16::from_be_bytes(itu_t35_bytes[1..3].try_into()?);
+                let itu_t_t35_terminal_provider_oriented_code =
+                    u16::from_be_bytes(itu_t35_bytes[3..5].try_into()?);
+
+                if itu_t_t35_country_code == 0xB5
+                    && itu_t_t35_terminal_provider_code == 0x003C
+                    && itu_t_t35_terminal_provider_oriented_code == 0x0001
+                {
+                    let application_identifier = itu_t35_bytes[5];
+                    let application_version = itu_t35_bytes[6];
+
+                    if application_identifier == 4 && application_version == 1 {
+                        return Ok(true);
+                    }
+                }
+            }
+        }
+    }
+
+    Ok(false)
+}
diff --git a/src/core/parser.rs b/src/core/parser.rs
@@ -1,19 +1,18 @@
-use std::convert::TryInto;
 use std::io::{stdout, BufRead, BufReader, BufWriter, Write};
 use std::path::PathBuf;
 use std::{fs::File, path::Path};
 
 use anyhow::{bail, Result};
 use indicatif::ProgressBar;
 
-use hevc_parser::hevc::{Frame, NALUnit, SeiMessage};
-use hevc_parser::hevc::{NAL_SEI_PREFIX, USER_DATA_REGISTERED_ITU_T_35};
+use hevc_parser::hevc::NAL_SEI_PREFIX;
+use hevc_parser::hevc::{Frame, NALUnit};
 use hevc_parser::HevcParser;
 
 use hdr10plus::metadata::Hdr10PlusMetadata;
 use hdr10plus::metadata_json::generate_json;
 
-use super::Format;
+use super::{is_st2094_40_sei, Format};
 
 pub const TOOL_NAME: &str = env!("CARGO_PKG_NAME");
 pub const TOOL_VERSION: &str = env!("CARGO_PKG_VERSION");
@@ -203,33 +202,8 @@ impl Parser {
             if let NAL_SEI_PREFIX = nal.nal_type {
                 let sei_payload = &data[nal.start..nal.end];
 
-                if sei_payload.len() >= 4 {
-                    let sei = SeiMessage::from_bytes(sei_payload)?;
-
-                    if sei.payload_type == USER_DATA_REGISTERED_ITU_T_35 {
-                        // FIXME: Not sure why 4 bytes..
-                        let itu_t35_bytes = &sei_payload[4..];
-
-                        if itu_t35_bytes.len() >= 7 {
-                            let itu_t_t35_country_code = itu_t35_bytes[0];
-                            let itu_t_t35_terminal_provider_code =
-                                u16::from_be_bytes(itu_t35_bytes[1..3].try_into()?);
-                            let itu_t_t35_terminal_provider_oriented_code =
-                                u16::from_be_bytes(itu_t35_bytes[3..5].try_into()?);
-
-                            if itu_t_t35_country_code == 0xB5
-                                && itu_t_t35_terminal_provider_code == 0x003C
-                                && itu_t_t35_terminal_provider_oriented_code == 0x0001
-                            {
-                                let application_identifier = itu_t35_bytes[5];
-                                let application_version = itu_t35_bytes[6];
-
-                                if application_identifier == 4 && application_version == 1 {
-                                    found_list.push(itu_t35_bytes.to_vec());
-                                }
-                            }
-                        }
-                    }
+                if is_st2094_40_sei(sei_payload)? {
+                    found_list.push(sei_payload[4..].to_vec());
                 }
             }
         }