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
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151

//!
//! Implement the muxer trait from av-format and expose all the correct
//! abstraction to handle them. Refer to the `Muxer` trait for more info.
//!
//!

use crate::common::Codec;
use av_data::packet::Packet;
use av_data::params::MediaKind;
use av_data::value::Value;
use av_format::common::GlobalInfo;
use av_format::error::*;
use av_format::muxer::Muxer;
use std::sync::Arc;

use av_bitstream::bytewrite::*;
use std::io::Write;

#[derive(Default, Debug)]
pub struct IvfMuxer {
    version: u16,
    width: u16,
    height: u16,
    rate: u32,
    scale: u32,
    codec: Codec,
    duration: u32,
    info: Option<GlobalInfo>,
}

impl IvfMuxer {
    pub fn new() -> IvfMuxer {
        IvfMuxer::default()
    }
}

/// This should be called if IvfMuxer::info is set
impl Muxer for IvfMuxer {
    fn configure(&mut self) -> Result<()> {
        match self.info.as_ref() {
            Some(info) if !info.streams.is_empty() => {
                self.duration = info.streams[0].duration.unwrap_or_default() as u32;
                let params = &info.streams[0].params;
                self.version = 0;
                if let Some(MediaKind::Video(video)) = &params.kind {
                    self.width = video.width as u16;
                    self.height = video.height as u16;
                };
                // TODO: parse scale
                self.rate = params.bit_rate as u32;
                self.scale = 1;
                self.codec = match params.codec_id.as_deref() {
                    Some("av1") => Codec::AV1,
                    Some("vp8") => Codec::VP8,
                    Some("vp9") => Codec::VP9,
                    _ => Codec::default(),
                };

                debug!("Configuration changes {:?}", self);

                Ok(())
            }

            _ => {
                debug!("No configuration changes {:?}", self);
                Ok(())
            }
        }
    }

    fn write_header(&mut self, buf: &mut Vec<u8>) -> Result<()> {
        debug!("Write muxer header: {:?}", self);

        buf.resize(32, 0);

        let codec = match self.codec {
            Codec::VP8 => b"VP80",
            Codec::VP9 => b"VP90",
            Codec::AV1 => b"AV01",
        };

        (&mut buf[0..=3]).write_all(b"DKIF")?;
        put_u16l(&mut buf[4..=5], self.version);
        put_u16l(&mut buf[6..=7], 32);
        (&mut buf[8..=11]).write_all(codec)?;
        put_u16l(&mut buf[12..=13], self.width);
        put_u16l(&mut buf[14..=15], self.height);
        put_u32l(&mut buf[16..=19], self.rate);
        put_u32l(&mut buf[20..=23], self.scale);
        put_u32l(&mut buf[24..=27], self.duration);
        put_u32l(&mut buf[28..=31], 0);

        Ok(())
    }

    fn write_packet(&mut self, buf: &mut Vec<u8>, pkt: Arc<Packet>) -> Result<()> {
        trace!("Write packet: {:?}", pkt.pos);

        let mut frame_header = [0; 12];

        put_u32l(&mut frame_header[0..4], pkt.data.len() as u32);
        put_u64l(&mut frame_header[4..12], pkt.pos.unwrap_or_default() as u64);

        buf.extend(&frame_header);
        buf.extend(&pkt.data);

        Ok(())
    }

    fn write_trailer(&mut self, _buf: &mut Vec<u8>) -> Result<()> {
        Ok(())
    }

    fn set_global_info(&mut self, info: GlobalInfo) -> Result<()> {
        self.info = Some(info);
        Ok(())
    }

    fn set_option<'a>(&mut self, _key: &str, _val: Value<'a>) -> Result<()> {
        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use av_format::common::GlobalInfo;
    use av_format::muxer::Context;
    use std::fs::File;

    #[test]
    fn mux() {
        let _ = pretty_env_logger::try_init();

        let output: File = tempfile::tempfile().unwrap();

        let info = GlobalInfo {
            duration: None,
            timebase: None,
            streams: Vec::new(),
        };

        let mux = Box::new(IvfMuxer::new());

        let mut muxer = Context::new(mux, Box::new(output));

        muxer.set_global_info(info).unwrap();
        muxer.configure().unwrap();
        muxer.write_header().unwrap();
    }
}