diff --git a/Cargo.toml b/Cargo.toml index 1396502..4bffc16 100644 --- a/Cargo.toml +++ b/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"] diff --git a/src/hevc/mod.rs b/src/hevc/mod.rs index 23b7594..89547b7 100644 --- a/src/hevc/mod.rs +++ b/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, diff --git a/src/hevc/pps.rs b/src/hevc/pps.rs index 5484cf8..2428eb0 100644 --- a/src/hevc/pps.rs +++ b/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, diff --git a/src/hevc/sps.rs b/src/hevc/sps.rs index 65a24cb..3b66c71 100644 --- a/src/hevc/sps.rs +++ b/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, diff --git a/src/hevc/vps.rs b/src/hevc/vps.rs index 9a00fac..353d251 100644 --- a/src/hevc/vps.rs +++ b/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, diff --git a/src/io/mod.rs b/src/io/mod.rs new file mode 100644 index 0000000..209c719 --- /dev/null +++ b/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, +} + +/// 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, +} + +pub fn format_from_path(input: &Path) -> Result { + 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"), + } + } +} diff --git a/src/io/processor.rs b/src/io/processor.rs new file mode 100644 index 0000000..1afbfca --- /dev/null +++ b/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, + sec_buf: Vec, + consumed: usize, + + chunk: Vec, + end: Vec, + offsets: Vec, + + 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(()) + } +} diff --git a/src/lib.rs b/src/lib.rs index 4cb2d19..852242d 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -6,6 +6,9 @@ use bitvec_helpers::bitvec_reader::BitVecReader; pub mod hevc; pub mod utils; +#[cfg(feature = "hevc_io")] +pub mod io; + use hevc::*; use pps::PPSNAL; use slice::SliceNAL; @@ -22,6 +25,7 @@ const HEADER_LEN_4: usize = 4; const NAL_START_CODE_3: &[u8] = &[0, 0, 1]; const NAL_START_CODE_4: &[u8] = &[0, 0, 0, 1]; +#[derive(Debug, Copy, Clone)] pub enum NALUStartCode { Length3, Length4, @@ -64,10 +68,7 @@ impl HevcParser { let mut consumed = 0; - let nal_start_tag = match &self.nalu_start_code { - NALUStartCode::Length3 => NAL_START_CODE_3, - NALUStartCode::Length4 => NAL_START_CODE_4, - }; + let nal_start_tag = self.nalu_start_code.slice(); loop { match Self::take_until_nal(nal_start_tag, &data[consumed..]) { @@ -78,10 +79,7 @@ impl HevcParser { offsets.push(consumed); // nom consumes the tag, so add it back - consumed += match &self.nalu_start_code { - NALUStartCode::Length3 => HEADER_LEN_3, - NALUStartCode::Length4 => HEADER_LEN_4, - }; + consumed += self.nalu_start_code.size(); } _ => return, } @@ -147,18 +145,23 @@ impl HevcParser { nal.end = end; nal.decoded_frame_index = self.decoded_index; - nal.start_code_len = if offset > 0 { + nal.start_code = if offset > 0 { // Previous byte is 0, offset..offset + 3 is [0, 0, 1] // Actual start code is length 4 if data[offset - 1] == 0 { - 4 + NALUStartCode::Length4 } else { - 3 + NALUStartCode::Length3 } } else { - 3 + NALUStartCode::Length3 }; + #[allow(deprecated)] + { + nal.start_code_len = nal.start_code.size() as u8; + } + if parse_nal { let bytes = clear_start_code_emulation_prevention_3_byte(&data[pos..parsing_end]); self.reader = BitVecReader::new(bytes); @@ -404,3 +407,19 @@ impl Default for NALUStartCode { NALUStartCode::Length3 } } + +impl NALUStartCode { + pub fn slice(&self) -> &[u8] { + match self { + NALUStartCode::Length3 => NAL_START_CODE_3, + NALUStartCode::Length4 => NAL_START_CODE_4, + } + } + + pub fn size(&self) -> usize { + match self { + NALUStartCode::Length3 => HEADER_LEN_3, + NALUStartCode::Length4 => HEADER_LEN_4, + } + } +}