Add HEVC processing helper

Author: quietvoid

Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "hevc_parser"
-version = "0.3.4"
+version = "0.4.0"
 authors = ["quietvoid"]
 edition = "2021"
 rust-version = "1.56.0"
@@ -11,4 +11,10 @@ repository = "https://github.com/quietvoid/hevc_parser"
 [dependencies]
 nom = "7.1.1"
 bitvec_helpers = "1.0.2"
-anyhow = "1.0.56"
+anyhow = "1.0.57"
+
+regex = { version = "1.5.5", optional = true }
+
+[features]
+default = ["hevc_io"]
+hevc_io = ["regex"]

src/hevc/mod.rs

@@ -2,7 +2,7 @@ use anyhow::{bail, Result};
 
 use self::slice::SliceNAL;
 
-use super::BitVecReader;
+use super::{BitVecReader, NALUStartCode};
 
 pub(crate) mod hrd_parameters;
 pub(crate) mod pps;
@@ -54,6 +54,10 @@ pub struct NALUnit {
     pub nal_type: u8,
     pub nuh_layer_id: u8,
     pub temporal_id: u8,
+
+    pub start_code: NALUStartCode,
+
+    #[deprecated(since = "0.4.0", note = "Please use `start_code` instead")]
     pub start_code_len: u8,
 
     pub decoded_frame_index: u64,

src/hevc/pps.rs

@@ -1,6 +1,7 @@
 use super::{scaling_list_data::ScalingListData, BitVecReader};
 use anyhow::Result;
 
+#[allow(clippy::upper_case_acronyms)]
 #[derive(Default, Debug, PartialEq)]
 pub struct PPSNAL {
     pub(crate) pps_id: u64,

src/hevc/sps.rs

@@ -6,6 +6,7 @@ use super::short_term_rps::ShortTermRPS;
 use super::vui_parameters::VuiParameters;
 use super::BitVecReader;
 
+#[allow(clippy::upper_case_acronyms)]
 #[derive(Default, Debug, PartialEq, Clone)]
 pub struct SPSNAL {
     pub(crate) vps_id: u8,

src/hevc/vps.rs

@@ -4,6 +4,7 @@ use super::hrd_parameters::HrdParameters;
 use super::profile_tier_level::ProfileTierLevel;
 use super::BitVecReader;
 
+#[allow(clippy::upper_case_acronyms)]
 #[derive(Default, Debug, PartialEq)]
 pub struct VPSNAL {
     pub(crate) vps_id: u8,

src/io/mod.rs

@@ -0,0 +1,83 @@
+use std::path::{Path, PathBuf};
+
+use anyhow::{bail, format_err, Result};
+use regex::Regex;
+
+pub mod processor;
+
+use super::{HevcParser, NALUStartCode, NALUnit};
+
+#[derive(Debug, PartialEq, Clone)]
+pub enum IoFormat {
+    Raw,
+    RawStdin,
+    Matroska,
+}
+
+pub trait IoProcessor {
+    /// Input path
+    fn input(&self) -> &PathBuf;
+    /// If the processor has a progress bar, this updates every megabyte read
+    fn update_progress(&mut self, delta: u64);
+
+    /// NALU processing callback
+    /// This is called after reading a 100kB chunk of the file
+    /// The resulting NALs are always complete and unique
+    ///
+    /// The data can be access through `chunk`, using the NAL start/end indices
+    fn process_nals(&mut self, parser: &HevcParser, nals: &[NALUnit], chunk: &[u8]) -> Result<()>;
+
+    /// Finalize callback, when the stream is done being read
+    /// Called at the end of `HevcProcessor::process_io`
+    fn finalize(&mut self, parser: &HevcParser) -> Result<()>;
+}
+
+/// Data for a frame, with its decoded index
+pub struct FrameBuffer {
+    pub frame_number: u64,
+    pub nals: Vec<NalBuffer>,
+}
+
+/// Data for a NALU, with type
+/// The data does not include the start code
+pub struct NalBuffer {
+    pub nal_type: u8,
+    pub start_code: NALUStartCode,
+    pub data: Vec<u8>,
+}
+
+pub fn format_from_path(input: &Path) -> Result<IoFormat> {
+    let regex = Regex::new(r"\.(hevc|.?265|mkv)")?;
+    let file_name = match input.file_name() {
+        Some(file_name) => file_name
+            .to_str()
+            .ok_or_else(|| format_err!("Invalid file name"))?,
+        None => "",
+    };
+
+    if file_name == "-" {
+        Ok(IoFormat::RawStdin)
+    } else if regex.is_match(file_name) && input.is_file() {
+        if file_name.ends_with(".mkv") {
+            Ok(IoFormat::Matroska)
+        } else {
+            Ok(IoFormat::Raw)
+        }
+    } else if file_name.is_empty() {
+        bail!("Missing input.")
+    } else if !input.is_file() {
+        bail!("Input file doesn't exist.")
+    } else {
+        bail!("Invalid input file type.")
+    }
+}
+
+impl std::fmt::Display for IoFormat {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match *self {
+            IoFormat::Matroska => write!(f, "Matroska file"),
+            IoFormat::Raw => write!(f, "HEVC file"),
+            IoFormat::RawStdin => write!(f, "HEVC pipe"),
+        }
+    }
+}

src/io/processor.rs

@@ -0,0 +1,173 @@
+use anyhow::{bail, Result};
+use std::io::Read;
+
+use super::{HevcParser, IoFormat, IoProcessor};
+
+/// Base HEVC stream processor
+pub struct HevcProcessor {
+    opts: HevcProcessorOpts,
+
+    format: IoFormat,
+    parser: HevcParser,
+
+    chunk_size: usize,
+
+    main_buf: Vec<u8>,
+    sec_buf: Vec<u8>,
+    consumed: usize,
+
+    chunk: Vec<u8>,
+    end: Vec<u8>,
+    offsets: Vec<usize>,
+
+    last_buffered_frame: u64,
+}
+
+/// Options for the processor
+#[derive(Default)]
+pub struct HevcProcessorOpts {
+    /// Buffer a frame when using `parse_nalus`.
+    /// This stops the stream reading as soon as a full frame has been parsed.
+    pub buffer_frame: bool,
+}
+
+impl HevcProcessor {
+    /// Initialize a HEVC stream processor
+    pub fn new(format: IoFormat, opts: HevcProcessorOpts, chunk_size: usize) -> Self {
+        let sec_buf = if format == IoFormat::RawStdin {
+            vec![0; 50_000]
+        } else {
+            Vec::new()
+        };
+
+        Self {
+            opts,
+            format,
+            parser: HevcParser::default(),
+
+            chunk_size,
+            main_buf: vec![0; chunk_size],
+            sec_buf,
+            consumed: 0,
+
+            chunk: Vec::with_capacity(chunk_size),
+            end: Vec::with_capacity(chunk_size),
+            offsets: Vec::with_capacity(2048),
+
+            last_buffered_frame: 0,
+        }
+    }
+
+    /// Fully parse the input stream
+    pub fn process_io(
+        &mut self,
+        reader: &mut dyn Read,
+        processor: &mut dyn IoProcessor,
+    ) -> Result<()> {
+        self.parse_nalus(reader, processor)?;
+
+        self.parser.finish();
+
+        processor.finalize(&self.parser)?;
+
+        Ok(())
+    }
+
+    /// Parse NALUs from the stream
+    /// Depending on the options, this either:
+    ///   - Loops the entire stream until EOF
+    ///   - Loops until a complete frame has been parsed
+    /// In both cases, the processor callback is called when a NALU payload is ready.
+    pub fn parse_nalus(
+        &mut self,
+        reader: &mut dyn Read,
+        processor: &mut dyn IoProcessor,
+    ) -> Result<()> {
+        while let Ok(n) = reader.read(&mut self.main_buf) {
+            let mut read_bytes = n;
+            if read_bytes == 0 && self.end.is_empty() && self.chunk.is_empty() {
+                break;
+            }
+
+            if self.format == IoFormat::RawStdin {
+                self.chunk.extend_from_slice(&self.main_buf[..read_bytes]);
+
+                loop {
+                    match reader.read(&mut self.sec_buf) {
+                        Ok(num) => {
+                            if num > 0 {
+                                read_bytes += num;
+
+                                self.chunk.extend_from_slice(&self.sec_buf[..num]);
+
+                                if read_bytes >= self.chunk_size {
+                                    break;
+                                }
+                            } else {
+                                break;
+                            }
+                        }
+                        Err(e) => bail!("{:?}", e),
+                    }
+                }
+            } else if read_bytes < self.chunk_size {
+                self.chunk.extend_from_slice(&self.main_buf[..read_bytes]);
+            } else {
+                self.chunk.extend_from_slice(&self.main_buf);
+            }
+
+            self.parser.get_offsets(&self.chunk, &mut self.offsets);
+
+            if self.offsets.is_empty() {
+                continue;
+            }
+
+            let last = if read_bytes < self.chunk_size {
+                *self.offsets.last().unwrap()
+            } else {
+                let last = self.offsets.pop().unwrap();
+
+                self.end.clear();
+                self.end.extend_from_slice(&self.chunk[last..]);
+
+                last
+            };
+
+            let nals = self
+                .parser
+                .split_nals(&self.chunk, &self.offsets, last, true)?;
+
+            // Process NALUs
+            processor.process_nals(&self.parser, &nals, &self.chunk)?;
+
+            self.chunk.clear();
+
+            if !self.end.is_empty() {
+                self.chunk.extend_from_slice(&self.end);
+                self.end.clear()
+            }
+
+            self.consumed += read_bytes;
+
+            if self.consumed >= 100_000_000 {
+                processor.update_progress(1);
+                self.consumed = 0;
+            }
+
+            if self.opts.buffer_frame {
+                let next_frame = nals.iter().map(|nal| nal.decoded_frame_index).max();
+
+                if let Some(number) = next_frame {
+                    if number > self.last_buffered_frame {
+                        self.last_buffered_frame = number;
+
+                        // Stop reading
+                        break;
+                    }
+                }
+            }
+        }
+
+        Ok(())
+    }
+}