webrtc_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::{Bytes, BytesMut};
8use rtp::packetizer::Depacketizer;
9
10use crate::error::Result;
11use crate::io::ivf_reader::IVFFileHeader;
12use crate::io::Writer;
13
14/// IVFWriter is used to take RTP packets and write them to an IVF on disk
15pub struct IVFWriter<W: Write + Seek> {
16    writer: W,
17    count: u64,
18    seen_key_frame: bool,
19    current_frame: Option<BytesMut>,
20    is_vp9: bool,
21}
22
23impl<W: Write + Seek> IVFWriter<W> {
24    /// new initialize a new IVF writer with an io.Writer output
25    pub fn new(writer: W, header: &IVFFileHeader) -> Result<Self> {
26        let mut w = IVFWriter {
27            writer,
28            count: 0,
29            seen_key_frame: false,
30            current_frame: None,
31            is_vp9: &header.four_cc != b"VP80",
32        };
33
34        w.write_header(header)?;
35
36        Ok(w)
37    }
38
39    fn write_header(&mut self, header: &IVFFileHeader) -> Result<()> {
40        self.writer.write_all(&header.signature)?; // DKIF
41        self.writer.write_u16::<LittleEndian>(header.version)?; // version
42        self.writer.write_u16::<LittleEndian>(header.header_size)?; // Header size
43        self.writer.write_all(&header.four_cc)?; // FOURCC
44        self.writer.write_u16::<LittleEndian>(header.width)?; // Width in pixels
45        self.writer.write_u16::<LittleEndian>(header.height)?; // Height in pixels
46        self.writer
47            .write_u32::<LittleEndian>(header.timebase_denominator)?; // Framerate denominator
48        self.writer
49            .write_u32::<LittleEndian>(header.timebase_numerator)?; // Framerate numerator
50        self.writer.write_u32::<LittleEndian>(header.num_frames)?; // Frame count, will be updated on first Close() call
51        self.writer.write_u32::<LittleEndian>(header.unused)?; // Unused
52
53        Ok(())
54    }
55}
56
57impl<W: Write + Seek> Writer for IVFWriter<W> {
58    /// write_rtp adds a new packet and writes the appropriate headers for it
59    fn write_rtp(&mut self, packet: &rtp::packet::Packet) -> Result<()> {
60        if packet.payload.is_empty() {
61            return Ok(());
62        }
63
64        let mut depacketizer: Box<dyn Depacketizer> = if self.is_vp9 {
65            Box::<rtp::codecs::vp9::Vp9Packet>::default()
66        } else {
67            Box::<rtp::codecs::vp8::Vp8Packet>::default()
68        };
69
70        let payload = depacketizer.depacketize(&packet.payload)?;
71
72        let is_key_frame = payload[0] & 0x01;
73
74        if (!self.seen_key_frame && is_key_frame == 1)
75            || (self.current_frame.is_none() && !depacketizer.is_partition_head(&packet.payload))
76        {
77            return Ok(());
78        }
79
80        self.seen_key_frame = true;
81        let frame_length = if let Some(current_frame) = &mut self.current_frame {
82            current_frame.extend(payload);
83            current_frame.len()
84        } else {
85            let mut current_frame = BytesMut::new();
86            current_frame.extend(payload);
87            let frame_length = current_frame.len();
88            self.current_frame = Some(current_frame);
89            frame_length
90        };
91
92        if !packet.header.marker {
93            return Ok(());
94        } else if let Some(current_frame) = &self.current_frame {
95            if current_frame.is_empty() {
96                return Ok(());
97            }
98        } else {
99            return Ok(());
100        }
101
102        self.writer.write_u32::<LittleEndian>(frame_length as u32)?; // Frame length
103        self.writer.write_u64::<LittleEndian>(self.count)?; // PTS
104        self.count += 1;
105
106        let frame_content = if let Some(current_frame) = self.current_frame.take() {
107            current_frame.freeze()
108        } else {
109            Bytes::new()
110        };
111
112        self.writer.write_all(&frame_content)?;
113
114        Ok(())
115    }
116
117    /// close stops the recording
118    fn close(&mut self) -> Result<()> {
119        // Update the frame count
120        self.writer.seek(SeekFrom::Start(24))?;
121        self.writer.write_u32::<LittleEndian>(self.count as u32)?;
122
123        self.writer.flush()?;
124        Ok(())
125    }
126}
webrtc_media/io/ivf_writer/mod.rs

webrtc_media/io/ivf_writer/
mod.rs
