diff --git a/src/mux/m2ts.rs b/src/mux/m2ts.rs index 5336565..ec4d6eb 100644 --- a/src/mux/m2ts.rs +++ b/src/mux/m2ts.rs @@ -35,6 +35,8 @@ pub struct M2tsStream { pending_frames: std::collections::VecDeque, pid_to_track: Vec<(u16, usize)>, pes_eof: bool, + /// Codec private data per stream (from FMKV header). + stored_codec_privates: Vec>>, } impl M2tsStream { @@ -53,6 +55,7 @@ impl M2tsStream { pending_frames: std::collections::VecDeque::new(), pid_to_track: Vec::new(), pes_eof: false, + stored_codec_privates: Vec::new(), } } @@ -91,6 +94,7 @@ impl M2tsStream { if let Ok(Some(m)) = meta::read_header(&mut reader) { let header_end = reader.stream_position()?; let content_size = file_size.saturating_sub(header_end); + let codec_privates = m.codec_privates(); let title = m.to_title(); let (pids, parsers, pid_to_track) = Self::setup_pes(&title.streams); return Ok(Self { @@ -105,6 +109,7 @@ impl M2tsStream { pending_frames: std::collections::VecDeque::new(), pid_to_track, pes_eof: false, + stored_codec_privates: codec_privates, }); } @@ -144,6 +149,7 @@ impl M2tsStream { pending_frames: std::collections::VecDeque::new(), pid_to_track, pes_eof: false, + stored_codec_privates: Vec::new(), }) } } @@ -197,6 +203,11 @@ impl crate::pes::Stream for M2tsStream { fn info(&self) -> &crate::disc::DiscTitle { &self.disc_title } fn codec_private(&self, track: usize) -> Option> { + // First check stored codec_privates from FMKV header + if let Some(Some(cp)) = self.stored_codec_privates.get(track) { + return Some(cp.clone()); + } + // Fall back to parser-extracted codec_private let pid = self.pid_to_track.iter() .find(|(_, idx)| *idx == track) .map(|(pid, _)| *pid)?; diff --git a/src/mux/meta.rs b/src/mux/meta.rs index 8f36913..f199c0e 100644 --- a/src/mux/meta.rs +++ b/src/mux/meta.rs @@ -46,6 +46,9 @@ pub enum MetaStream { label: String, #[serde(default)] secondary: bool, + /// Base64-encoded codec initialization data (HEVCDecoderConfigurationRecord, etc.) + #[serde(default, skip_serializing_if = "Option::is_none")] + codec_private: Option, }, #[serde(rename = "audio")] Audio { @@ -74,6 +77,20 @@ pub enum MetaStream { } impl M2tsMeta { + /// Build metadata from a disc Title with optional codec_private per stream. + pub fn from_title_with_privates(title: &DiscTitle, codec_privates: &[Option>]) -> Self { + let mut meta = Self::from_title(title); + for (i, s) in meta.streams.iter_mut().enumerate() { + if let MetaStream::Video { codec_private, .. } = s { + if let Some(Some(cp)) = codec_privates.get(i) { + use base64::Engine; + *codec_private = Some(base64::engine::general_purpose::STANDARD.encode(cp)); + } + } + } + meta + } + /// Build metadata from a disc Title. pub fn from_title(title: &DiscTitle) -> Self { let streams = title @@ -88,6 +105,7 @@ impl M2tsMeta { hdr: v.hdr.id().into(), label: v.label.clone(), secondary: v.secondary, + codec_private: None, }, Stream::Audio(a) => MetaStream::Audio { pid: a.pid, @@ -129,6 +147,7 @@ impl M2tsMeta { hdr, label, secondary, + codec_private: _, } => Stream::Video(VideoStream { pid: *pid, codec: codec.parse().unwrap(), @@ -183,6 +202,21 @@ impl M2tsMeta { content_format: crate::disc::ContentFormat::BdTs, } } + + /// Extract codec_private data per stream (from FMKV header). + /// Returns a Vec matching stream order — None for streams without codec_private. + pub fn codec_privates(&self) -> Vec>> { + self.streams.iter().map(|s| { + if let MetaStream::Video { codec_private: Some(ref b64), .. } = s { +{ + use base64::Engine; + base64::engine::general_purpose::STANDARD.decode(b64).ok() + } + } else { + None + } + }).collect() + } } /// Write the metadata header to a writer. Padded to 192-byte boundary. diff --git a/src/mux/mkvstream.rs b/src/mux/mkvstream.rs index 491fd86..90ac114 100644 --- a/src/mux/mkvstream.rs +++ b/src/mux/mkvstream.rs @@ -207,6 +207,22 @@ impl crate::pes::Stream for MkvStream { fn finish(&mut self) -> io::Result<()> { Ok(()) } fn info(&self) -> &crate::disc::DiscTitle { &self.disc_title } + + fn codec_private(&self, track: usize) -> Option> { + let track_num = (track + 1) as u16; // MKV tracks are 1-based + if let Mode::Read(ref rs) = self.mode { + rs.codec_privates + .iter() + .find(|(tn, _)| *tn == track_num) + .map(|(_, data)| data.clone()) + } else { + None + } + } + + fn headers_ready(&self) -> bool { + true // MKV has all headers upfront in the EBML header + } } impl IOStream for MkvStream { diff --git a/src/mux/pesout.rs b/src/mux/pesout.rs index 46ad158..4c7ae3c 100644 --- a/src/mux/pesout.rs +++ b/src/mux/pesout.rs @@ -13,17 +13,28 @@ pub struct M2tsOutputStream { } impl M2tsOutputStream { - pub fn create(path: &str, title: &DiscTitle) -> io::Result { + pub fn create( + path: &str, + title: &DiscTitle, + codec_privates: &[Option>], + ) -> io::Result { let file = std::fs::File::create(path) .map_err(|e| io::Error::new(e.kind(), format!("m2ts://{}: {}", path, e)))?; let mut writer = io::BufWriter::with_capacity(4 * 1024 * 1024, file); - // Write FMKV metadata header so M2tsStream::open can read it back + // Write FMKV metadata header with codec_privates so M2tsStream::open can read them back if !title.streams.is_empty() { - let m = super::meta::M2tsMeta::from_title(title); + let m = super::meta::M2tsMeta::from_title_with_privates(title, codec_privates); super::meta::write_header(&mut writer, &m)?; } let pids = extract_pids(title); - Ok(Self { muxer: TsMuxer::new(writer, &pids), title: title.clone() }) + let mut muxer = TsMuxer::new(writer, &pids); + // Pass codec_privates to TsMuxer for Annex B parameter set injection + for (i, cp) in codec_privates.iter().enumerate() { + if let Some(data) = cp { + muxer.set_codec_private(i, data.clone()); + } + } + Ok(Self { muxer, title: title.clone() }) } } diff --git a/src/mux/resolve.rs b/src/mux/resolve.rs index 5946187..bae4b21 100644 --- a/src/mux/resolve.rs +++ b/src/mux/resolve.rs @@ -384,7 +384,7 @@ pub fn output( } StreamUrl::M2ts { ref path } => { validate_file_path(path, "m2ts")?; - Ok(Box::new(super::pesout::M2tsOutputStream::create(&path.to_string_lossy(), title)?)) + Ok(Box::new(super::pesout::M2tsOutputStream::create(&path.to_string_lossy(), title, codec_privates)?)) } StreamUrl::Network { ref addr } => { validate_network_addr(addr)?; diff --git a/src/mux/tsmux.rs b/src/mux/tsmux.rs index 412bd4d..5aa53ed 100644 --- a/src/mux/tsmux.rs +++ b/src/mux/tsmux.rs @@ -12,20 +12,34 @@ const TS_PAYLOAD: usize = 184; pub struct TsMuxer { writer: W, pids: Vec, - continuity: Vec, // per-PID continuity counter (0-15) + continuity: Vec, // per-PID continuity counter (0-15) + codec_privates: Vec>>, // per-track codec_private (for video parameter sets) + params_written: Vec, // per-track: have we written parameter sets? } impl TsMuxer { pub fn new(writer: W, pids: &[u16]) -> Self { - let continuity = vec![0u8; pids.len()]; + let n = pids.len(); Self { writer, pids: pids.to_vec(), - continuity, + continuity: vec![0u8; n], + codec_privates: vec![None; n], + params_written: vec![false; n], + } + } + + /// Set codec_private data for a track. Used to prepend VPS/SPS/PPS + /// as Annex B NALs before the first keyframe in the transport stream. + pub fn set_codec_private(&mut self, track: usize, data: Vec) { + if track < self.codec_privates.len() { + self.codec_privates[track] = Some(data); } } /// Write a PES frame as BD-TS packets. + /// Video frame data is expected as length-prefixed NALUs (MKV/PES format) + /// and is converted to Annex B for transport stream. pub fn write_frame( &mut self, track: usize, @@ -36,11 +50,31 @@ impl TsMuxer { return Ok(()); // unknown track, skip } let pid = self.pids[track]; + let is_video = (0x1011..=0x101F).contains(&pid); + + // For video: convert length-prefixed NALUs to Annex B (start codes) + // On first keyframe, prepend parameter sets from codec_private + let es_data = if is_video && !data.is_empty() { + let mut annex_b = Vec::new(); + // Prepend codec_private parameter sets on first keyframe + if !self.params_written[track] { + if let Some(ref cp) = self.codec_privates[track] { + if let Some(params) = hvcc_to_annex_b(cp) { + annex_b.extend_from_slice(¶ms); + self.params_written[track] = true; + } + } + } + annex_b.extend_from_slice(&length_prefixed_to_annex_b(data)); + annex_b + } else { + data.to_vec() + }; // Build PES packet: header + data let pts_90k = (pts_ns * 9 / 100_000) as u64; - let pes_header = build_pes_header(pid, pts_90k, data.len()); - let pes_packet = [&pes_header[..], data].concat(); + let pes_header = build_pes_header(pid, pts_90k, es_data.len()); + let pes_packet = [&pes_header[..], &es_data[..]].concat(); // Split into TS packets let mut offset = 0; @@ -152,3 +186,63 @@ fn build_pes_header(pid: u16, pts_90k: u64, data_len: usize) -> Vec { header } + +/// Extract NAL arrays from HEVCDecoderConfigurationRecord and convert to Annex B. +/// Returns VPS + SPS + PPS as Annex B NAL units (00 00 00 01 + NAL). +fn hvcc_to_annex_b(hvcc: &[u8]) -> Option> { + // HEVCDecoderConfigurationRecord: 22 bytes header, then NAL arrays + if hvcc.len() < 23 { + return None; + } + let num_arrays = hvcc[22] as usize; + let mut out = Vec::new(); + let mut offset = 23; + + for _ in 0..num_arrays { + if offset + 3 > hvcc.len() { break; } + // array: 1 byte (completeness + NAL type), 2 bytes (numNalus) + let _nal_type = hvcc[offset] & 0x3F; + let num_nalus = u16::from_be_bytes([hvcc[offset + 1], hvcc[offset + 2]]) as usize; + offset += 3; + + for _ in 0..num_nalus { + if offset + 2 > hvcc.len() { break; } + let nal_len = u16::from_be_bytes([hvcc[offset], hvcc[offset + 1]]) as usize; + offset += 2; + if offset + nal_len > hvcc.len() { break; } + out.extend_from_slice(&[0x00, 0x00, 0x00, 0x01]); + out.extend_from_slice(&hvcc[offset..offset + nal_len]); + offset += nal_len; + } + } + + if out.is_empty() { None } else { Some(out) } +} + +/// Convert length-prefixed NALUs (4-byte BE length + NAL) to Annex B +/// (00 00 00 01 + NAL). Used for video elementary streams in TS. +fn length_prefixed_to_annex_b(data: &[u8]) -> Vec { + let mut out = Vec::with_capacity(data.len()); + let mut offset = 0; + while offset + 4 <= data.len() { + let len = u32::from_be_bytes([ + data[offset], + data[offset + 1], + data[offset + 2], + data[offset + 3], + ]) as usize; + offset += 4; + if offset + len > data.len() { + break; + } + out.extend_from_slice(&[0x00, 0x00, 0x00, 0x01]); + out.extend_from_slice(&data[offset..offset + len]); + offset += len; + } + // If data doesn't look like length-prefixed NALs (no valid parse), + // return original data unchanged — it may already be Annex B. + if out.is_empty() && !data.is_empty() { + return data.to_vec(); + } + out +}