1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113

#[cfg(test)]
mod ivf_writer_test;

use crate::io::ivf_reader::IVFFileHeader;
use crate::io::Writer;
use anyhow::Result;
use byteorder::{LittleEndian, WriteBytesExt};
use bytes::{Bytes, BytesMut};
use rtp::packetizer::Depacketizer;
use std::io::{Seek, SeekFrom, Write};

/// IVFWriter is used to take RTP packets and write them to an IVF on disk
pub struct IVFWriter<W: Write + Seek> {
    writer: W,
    count: u64,
    seen_key_frame: bool,
    current_frame: Option<BytesMut>,
}

impl<W: Write + Seek> IVFWriter<W> {
    /// new initialize a new IVF writer with an io.Writer output
    pub fn new(writer: W, header: &IVFFileHeader) -> Result<Self> {
        let mut w = IVFWriter {
            writer,
            count: 0,
            seen_key_frame: false,
            current_frame: None,
        };

        w.write_header(header)?;

        Ok(w)
    }

    fn write_header(&mut self, header: &IVFFileHeader) -> Result<()> {
        self.writer.write_all(&header.signature)?; // DKIF
        self.writer.write_u16::<LittleEndian>(header.version)?; // version
        self.writer.write_u16::<LittleEndian>(header.header_size)?; // Header size
        self.writer.write_all(&header.four_cc)?; // FOURCC
        self.writer.write_u16::<LittleEndian>(header.width)?; // Width in pixels
        self.writer.write_u16::<LittleEndian>(header.height)?; // Height in pixels
        self.writer
            .write_u32::<LittleEndian>(header.timebase_denominator)?; // Framerate denominator
        self.writer
            .write_u32::<LittleEndian>(header.timebase_numerator)?; // Framerate numerator
        self.writer.write_u32::<LittleEndian>(header.num_frames)?; // Frame count, will be updated on first Close() call
        self.writer.write_u32::<LittleEndian>(header.unused)?; // Unused

        Ok(())
    }
}

impl<W: Write + Seek> Writer for IVFWriter<W> {
    /// write_rtp adds a new packet and writes the appropriate headers for it
    fn write_rtp(&mut self, packet: &rtp::packet::Packet) -> Result<()> {
        let mut vp8packet = rtp::codecs::vp8::Vp8Packet::default();
        vp8packet.depacketize(&packet.payload)?;

        let is_key_frame = vp8packet.payload[0] & 0x01;

        if (!self.seen_key_frame && is_key_frame == 1)
            || (self.current_frame.is_none() && vp8packet.s != 1)
        {
            return Ok(());
        }

        self.seen_key_frame = true;
        let frame_length = if let Some(current_frame) = &mut self.current_frame {
            current_frame.extend(vp8packet.payload);
            current_frame.len()
        } else {
            let mut current_frame = BytesMut::new();
            current_frame.extend(vp8packet.payload);
            let frame_length = current_frame.len();
            self.current_frame = Some(current_frame);
            frame_length
        };

        if !packet.header.marker {
            return Ok(());
        } else if let Some(current_frame) = &self.current_frame {
            if current_frame.is_empty() {
                return Ok(());
            }
        } else {
            return Ok(());
        }

        self.writer.write_u32::<LittleEndian>(frame_length as u32)?; // Frame length
        self.writer.write_u64::<LittleEndian>(self.count)?; // PTS
        self.count += 1;

        let frame_content = if let Some(current_frame) = self.current_frame.take() {
            current_frame.freeze()
        } else {
            Bytes::new()
        };

        self.writer.write_all(&frame_content)?;

        Ok(())
    }

    /// close stops the recording
    fn close(&mut self) -> Result<()> {
        // Update the frame count
        self.writer.seek(SeekFrom::Start(24))?;
        self.writer.write_u32::<LittleEndian>(self.count as u32)?;

        self.writer.flush()?;
        Ok(())
    }
}