const AV2_FOURCC: u32 = 0x3230_5641;
pub struct IvfWriter {
buf: Vec<u8>,
frames: u32,
count_off: usize,
}
impl IvfWriter {
pub fn new(width: u16, height: u16, fps_num: u32, fps_den: u32) -> Self {
let mut buf = Vec::with_capacity(64);
buf.extend_from_slice(b"DKIF");
buf.extend_from_slice(&0u16.to_le_bytes()); buf.extend_from_slice(&32u16.to_le_bytes()); buf.extend_from_slice(&AV2_FOURCC.to_le_bytes());
buf.extend_from_slice(&width.to_le_bytes());
buf.extend_from_slice(&height.to_le_bytes());
buf.extend_from_slice(&fps_num.to_le_bytes()); buf.extend_from_slice(&fps_den.to_le_bytes()); let count_off = buf.len();
buf.extend_from_slice(&0u32.to_le_bytes()); buf.extend_from_slice(&0u32.to_le_bytes()); Self {
buf,
frames: 0,
count_off,
}
}
pub fn write_frame(&mut self, data: &[u8], pts: u64) {
self.buf
.extend_from_slice(&(data.len() as u32).to_le_bytes());
self.buf
.extend_from_slice(&((pts & 0xFFFF_FFFF) as u32).to_le_bytes());
self.buf
.extend_from_slice(&((pts >> 32) as u32).to_le_bytes());
self.buf.extend_from_slice(data);
self.frames += 1;
}
pub fn finish(mut self) -> Vec<u8> {
self.buf[self.count_off..self.count_off + 4].copy_from_slice(&self.frames.to_le_bytes());
self.buf
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn header_is_32_bytes_and_fourcc_av02() {
let w = IvfWriter::new(320, 240, 30, 1);
let out = w.finish();
assert_eq!(&out[0..4], b"DKIF");
assert_eq!(u16::from_le_bytes([out[6], out[7]]), 32);
assert_eq!(&out[8..12], b"AV02");
assert_eq!(u16::from_le_bytes([out[12], out[13]]), 320);
assert_eq!(u32::from_le_bytes([out[24], out[25], out[26], out[27]]), 0);
}
#[test]
fn frame_count_patched() {
let mut w = IvfWriter::new(16, 16, 30, 1);
w.write_frame(&[1, 2, 3], 0);
w.write_frame(&[4, 5], 1);
let out = w.finish();
assert_eq!(u32::from_le_bytes([out[24], out[25], out[26], out[27]]), 2);
assert_eq!(u32::from_le_bytes([out[32], out[33], out[34], out[35]]), 3);
}
}