Fix M2TS/MKV roundtrip: codec_private in FMKV header, Annex B conversion
Bug 1: M2TS roundtrip dropped frames — TsMuxer converts length-prefixed NALs to Annex B, prepends VPS/SPS/PPS from HEVCDecoderConfigurationRecord. Bug 2: MKV remux lost codec_private — MkvStream.codec_private() now returns data from EBML header. FMKV header carries codec_private (base64) per video stream for lossless M2TS roundtrip.
This commit is contained in:
+99
-5
@@ -12,20 +12,34 @@ const TS_PAYLOAD: usize = 184;
|
||||
pub struct TsMuxer<W: Write> {
|
||||
writer: W,
|
||||
pids: Vec<u16>,
|
||||
continuity: Vec<u8>, // per-PID continuity counter (0-15)
|
||||
continuity: Vec<u8>, // per-PID continuity counter (0-15)
|
||||
codec_privates: Vec<Option<Vec<u8>>>, // per-track codec_private (for video parameter sets)
|
||||
params_written: Vec<bool>, // per-track: have we written parameter sets?
|
||||
}
|
||||
|
||||
impl<W: Write> TsMuxer<W> {
|
||||
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<u8>) {
|
||||
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<W: Write> TsMuxer<W> {
|
||||
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<u8> {
|
||||
|
||||
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<Vec<u8>> {
|
||||
// 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<u8> {
|
||||
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
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user