rtc_media/io/ivf_writer/
mod.rs

1#[cfg(test)]
2mod ivf_writer_test;
3
4use std::io::{Seek, SeekFrom, Write};
5
6use byteorder::{LittleEndian, WriteBytesExt};
7use bytes::{BufMut, Bytes, BytesMut};
8use rtp::codec::av1::Av1Depacketizer;
9use rtp::packetizer::Depacketizer;
10
11use crate::io::Writer;
12use crate::io::ivf_reader::IVFFileHeader;
13use shared::error::Result;
14
15/// Codec type for IVF writer
16#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
17pub enum IvfCodec {
18    #[default]
19    Vp8,
20    Vp9,
21    Av1,
22}
23
24impl IvfCodec {
25    /// Create codec from FOURCC bytes
26    pub fn from_fourcc(fourcc: &[u8; 4]) -> Self {
27        match fourcc {
28            b"VP80" => IvfCodec::Vp8,
29            b"VP90" => IvfCodec::Vp9,
30            b"AV01" => IvfCodec::Av1,
31            _ => IvfCodec::Vp8, // Default fallback
32        }
33    }
34
35    /// Get FOURCC bytes for this codec
36    pub fn fourcc(&self) -> [u8; 4] {
37        match self {
38            IvfCodec::Vp8 => *b"VP80",
39            IvfCodec::Vp9 => *b"VP90",
40            IvfCodec::Av1 => *b"AV01",
41        }
42    }
43}
44
45/// IVFWriter is used to take RTP packets and write them to an IVF on disk
46pub struct IVFWriter<W: Write + Seek> {
47    writer: W,
48    count: u64,
49    seen_key_frame: bool,
50    current_frame: Option<BytesMut>,
51    codec: IvfCodec,
52    /// AV1 depacketizer (lazily initialized)
53    av1_depacketizer: Option<Av1Depacketizer>,
54}
55
56impl<W: Write + Seek> IVFWriter<W> {
57    /// new initialize a new IVF writer with an io.Writer output
58    pub fn new(writer: W, header: &IVFFileHeader) -> Result<Self> {
59        let codec = IvfCodec::from_fourcc(&header.four_cc);
60
61        let mut w = IVFWriter {
62            writer,
63            count: 0,
64            seen_key_frame: false,
65            current_frame: None,
66            codec,
67            av1_depacketizer: None,
68        };
69
70        w.write_header(header)?;
71
72        Ok(w)
73    }
74
75    fn write_header(&mut self, header: &IVFFileHeader) -> Result<()> {
76        self.writer.write_all(&header.signature)?; // DKIF
77        self.writer.write_u16::<LittleEndian>(header.version)?; // version
78        self.writer.write_u16::<LittleEndian>(header.header_size)?; // Header size
79        self.writer.write_all(&header.four_cc)?; // FOURCC
80        self.writer.write_u16::<LittleEndian>(header.width)?; // Width in pixels
81        self.writer.write_u16::<LittleEndian>(header.height)?; // Height in pixels
82        self.writer
83            .write_u32::<LittleEndian>(header.timebase_denominator)?; // Framerate denominator
84        self.writer
85            .write_u32::<LittleEndian>(header.timebase_numerator)?; // Framerate numerator
86        self.writer.write_u32::<LittleEndian>(header.num_frames)?; // Frame count, will be updated on first Close() call
87        self.writer.write_u32::<LittleEndian>(header.unused)?; // Unused
88
89        Ok(())
90    }
91
92    /// Write VP8 packet
93    fn write_vp8(&mut self, packet: &rtp::Packet) -> Result<()> {
94        let mut depacketizer = rtp::codec::vp8::Vp8Packet::default();
95        let payload = depacketizer.depacketize(&packet.payload)?;
96
97        // VP8 keyframe: first bit of first byte is 0
98        let is_key_frame = (payload[0] & 0x01) == 0;
99
100        if !self.seen_key_frame && !is_key_frame {
101            return Ok(());
102        }
103        if self.current_frame.is_none() && !depacketizer.is_partition_head(&packet.payload) {
104            return Ok(());
105        }
106
107        self.seen_key_frame = true;
108        self.append_to_frame(payload);
109
110        if !packet.header.marker {
111            return Ok(());
112        }
113
114        self.write_current_frame()
115    }
116
117    /// Write VP9 packet
118    fn write_vp9(&mut self, packet: &rtp::Packet) -> Result<()> {
119        let mut depacketizer = rtp::codec::vp9::Vp9Packet::default();
120        let payload = depacketizer.depacketize(&packet.payload)?;
121
122        // VP9 keyframe: P bit is 0 (inter-picture predicted frame = false)
123        let is_key_frame = !depacketizer.p;
124
125        if !self.seen_key_frame && !is_key_frame {
126            return Ok(());
127        }
128        if self.current_frame.is_none() && !depacketizer.b {
129            return Ok(());
130        }
131
132        self.seen_key_frame = true;
133        self.append_to_frame(payload);
134
135        if !packet.header.marker {
136            return Ok(());
137        }
138
139        self.write_current_frame()
140    }
141
142    /// Write AV1 packet
143    fn write_av1(&mut self, packet: &rtp::Packet) -> Result<()> {
144        // Initialize depacketizer if needed
145        if self.av1_depacketizer.is_none() {
146            self.av1_depacketizer = Some(Av1Depacketizer::new());
147        }
148
149        let depacketizer = self.av1_depacketizer.as_mut().unwrap();
150        let payload = depacketizer.depacketize(&packet.payload)?;
151
152        if !self.seen_key_frame {
153            // AV1 keyframe: N bit set or first OBU is sequence header
154            // OBU type for sequence header = 1, located at bits 3-6 of first byte
155            let is_key_frame =
156                depacketizer.n || (!payload.is_empty() && ((payload[0] & 0x78) >> 3) == 1);
157            if !is_key_frame {
158                return Ok(());
159            }
160            self.seen_key_frame = true;
161        }
162
163        self.append_to_frame(payload);
164
165        if !packet.header.marker {
166            return Ok(());
167        }
168
169        // For AV1, prepend temporal delimiter OBU before each frame
170        // Temporal delimiter: type=2, has_size=1, size=0
171        // OBU header: 0b0001_0010 = 0x12 (type=2, has_size_field=1)
172        let mut frame_with_delimiter = BytesMut::with_capacity(2 + self.current_frame_len());
173        frame_with_delimiter.put_u8(0x12); // OBU header: temporal delimiter with size field
174        frame_with_delimiter.put_u8(0x00); // Size = 0
175
176        if let Some(current_frame) = self.current_frame.take() {
177            frame_with_delimiter.extend_from_slice(&current_frame);
178        }
179
180        self.write_frame(&frame_with_delimiter)
181    }
182
183    /// Append payload to current frame
184    fn append_to_frame(&mut self, payload: Bytes) {
185        if let Some(current_frame) = &mut self.current_frame {
186            current_frame.extend(payload);
187        } else {
188            let mut current_frame = BytesMut::new();
189            current_frame.extend(payload);
190            self.current_frame = Some(current_frame);
191        }
192    }
193
194    /// Get current frame length
195    fn current_frame_len(&self) -> usize {
196        self.current_frame.as_ref().map_or(0, |f| f.len())
197    }
198
199    /// Write current frame to output
200    fn write_current_frame(&mut self) -> Result<()> {
201        if let Some(current_frame) = &self.current_frame {
202            if current_frame.is_empty() {
203                return Ok(());
204            }
205        } else {
206            return Ok(());
207        }
208
209        let frame_content = self.current_frame.take().unwrap();
210        self.write_frame(&frame_content)
211    }
212
213    /// Write frame data to output
214    fn write_frame(&mut self, frame: &[u8]) -> Result<()> {
215        self.writer.write_u32::<LittleEndian>(frame.len() as u32)?; // Frame length
216        self.writer.write_u64::<LittleEndian>(self.count)?; // PTS
217        self.count += 1;
218        self.writer.write_all(frame)?;
219        Ok(())
220    }
221}
222
223impl<W: Write + Seek> Writer for IVFWriter<W> {
224    /// write_rtp adds a new packet and writes the appropriate headers for it
225    fn write_rtp(&mut self, packet: &rtp::Packet) -> Result<()> {
226        if packet.payload.is_empty() {
227            return Ok(());
228        }
229
230        match self.codec {
231            IvfCodec::Vp8 => self.write_vp8(packet),
232            IvfCodec::Vp9 => self.write_vp9(packet),
233            IvfCodec::Av1 => self.write_av1(packet),
234        }
235    }
236
237    /// close stops the recording
238    fn close(&mut self) -> Result<()> {
239        // Update the frame count
240        self.writer.seek(SeekFrom::Start(24))?;
241        self.writer.write_u32::<LittleEndian>(self.count as u32)?;
242
243        self.writer.flush()?;
244        Ok(())
245    }
246}
rtc_media/io/ivf_writer/mod.rs

rtc_media/io/ivf_writer/
mod.rs
