]> git.nihav.org Git - nihav.git/commitdiff
avi muxer
authorKostya Shishkov <kostya.shishkov@gmail.com>
Thu, 21 May 2020 09:01:00 +0000 (11:01 +0200)
committerKostya Shishkov <kostya.shishkov@gmail.com>
Thu, 21 May 2020 09:01:00 +0000 (11:01 +0200)
nihav-commonfmt/Cargo.toml
nihav-commonfmt/src/lib.rs
nihav-commonfmt/src/muxers/avi.rs [new file with mode: 0644]
nihav-commonfmt/src/muxers/mod.rs [new file with mode: 0644]

index 6c0e81073ef9c4a38d9ffeb5c4c9377abbd4d4bd..5807da86ea19920da63bbe82259653641ca49021 100644 (file)
@@ -18,12 +18,15 @@ features = ["h263", "mdct", "fft", "dsp_window"]
 nihav_realmedia = { path = "../nihav-realmedia" }
 
 [features]
-default = ["all_decoders", "all_demuxers"]
+default = ["all_decoders", "all_demuxers", "all_muxers"]
 decoders = []
 demuxers = []
+muxers = []
 all_demuxers = ["demuxer_avi", "demuxer_mov"]
 demuxer_avi = ["demuxers"]
 demuxer_mov = ["demuxers"]
+all_muxers = ["muxer_avi"]
+muxer_avi = ["muxers"]
 
 all_decoders = ["all_video_decoders", "all_audio_decoders"]
 
index ef5ad072dcb6dd4c3a63ba8064d74fb1b611927a..82bfd493784b1a474b4c55ca1d3fd99728e80a31 100644 (file)
@@ -16,5 +16,10 @@ mod demuxers;
 #[cfg(feature="demuxers")]
 pub use crate::demuxers::generic_register_all_demuxers;
 
+#[cfg(feature="muxers")]
+mod muxers;
+#[cfg(feature="muxers")]
+pub use crate::muxers::generic_register_all_muxers;
+
 #[cfg(test)]
 extern crate nihav_realmedia;
diff --git a/nihav-commonfmt/src/muxers/avi.rs b/nihav-commonfmt/src/muxers/avi.rs
new file mode 100644 (file)
index 0000000..d47aac3
--- /dev/null
@@ -0,0 +1,339 @@
+use nihav_core::muxers::*;
+use nihav_registry::register::*;
+
+#[derive(Clone,Copy)]
+struct IdxEntry {
+    stream:     u32,
+    stype:      StreamType,
+    key:        bool,
+    pos:        u32,
+    len:        u32,
+}
+
+#[derive(Clone,Copy)]
+struct AVIStream {
+    strh_pos:   u64,
+    nframes:    u32,
+    is_video:   bool,
+    max_size:   u32,    
+}
+
+struct AVIMuxer<'a> {
+    bw:             &'a mut ByteWriter<'a>,
+    index:          Vec<IdxEntry>,
+    video_str:      Option<usize>,
+    video_id:       u32,
+    data_pos:       u64,
+    stream_info:    Vec<AVIStream>,
+}
+
+impl<'a> AVIMuxer<'a> {
+    fn new(bw: &'a mut ByteWriter<'a>) -> Self {
+        Self {
+            bw,
+            index:          Vec::new(),
+            video_str:      None,
+            video_id:       0,
+            data_pos:       0,
+            stream_info:    Vec::with_capacity(2),
+        }
+    }
+}
+
+fn patch_size(bw: &mut ByteWriter, pos: u64) -> MuxerResult<()> {
+    let size = bw.tell() - pos;
+    bw.seek(SeekFrom::Current(-((size + 4) as i64)))?;
+    bw.write_u32le(size as u32)?;
+    bw.seek(SeekFrom::End(0))?;
+    Ok(())
+}
+
+fn write_chunk_hdr(bw: &mut ByteWriter, stype: StreamType, str_no: u32) -> MuxerResult<()> {
+    bw.write_byte(b'0' + ((str_no / 10) as u8))?;
+    bw.write_byte(b'0' + ((str_no % 10) as u8))?;
+    match stype {
+        StreamType::Video => { bw.write_buf(b"dc")?; },
+        StreamType::Audio => { bw.write_buf(b"wb")?; },
+        StreamType::Subtitles => { bw.write_buf(b"tx")?; },
+        _ => return Err(MuxerError::UnsupportedFormat),
+    };
+    Ok(())
+}
+
+impl<'a> MuxCore<'a> for AVIMuxer<'a> {
+    fn create(&mut self, strmgr: &StreamManager) -> MuxerResult<()> {
+        if strmgr.get_num_streams() == 0 {
+            return Err(MuxerError::InvalidArgument);
+        }
+        if strmgr.get_num_streams() > 99 {
+            return Err(MuxerError::UnsupportedFormat);
+        }
+        for (str_no, str) in strmgr.iter().enumerate() {
+            if str.get_media_type() == StreamType::Video {
+                self.video_str = Some(str_no);
+                self.video_id  = str.id;
+                break;
+            }
+        }
+        let (vinfo, tb_num, tb_den) = if let Some(str_id) = self.video_str {
+                let vstr = strmgr.get_stream(str_id).unwrap();
+                (vstr.get_info(), vstr.tb_num, vstr.tb_den)
+            } else {
+                (NACodecInfo::new_dummy(), 0, 1)
+            };
+        let hdrl_pos = self.bw.tell() + 20;
+        self.bw.write_buf(b"RIFF\0\0\0\0AVI LIST\0\0\0\0hdrlavih")?;
+        self.bw.write_u32le(56)?; // avih size
+        let ms_per_frame = NATimeInfo::ts_to_time(1, 1000000, tb_num, tb_den);
+        self.bw.write_u32le(ms_per_frame as u32)?;
+        self.bw.write_u32le(0)?; // max transfer rate
+        self.bw.write_u32le(0)?; // padding granularity
+        self.bw.write_u32le(0)?; // flags
+        self.bw.write_u32le(0)?; // total frames
+        self.bw.write_u32le(0)?; // initial frames
+        self.bw.write_u32le(strmgr.get_num_streams() as u32)?;
+        self.bw.write_u32le(0)?; // suggested buffer size
+        if let NACodecTypeInfo::Video(ref vinfo) = vinfo.get_properties() {
+            self.bw.write_u32le(vinfo.width as u32)?;
+            self.bw.write_u32le(vinfo.height as u32)?;
+        } else {
+            self.bw.write_u32le(0)?;
+            self.bw.write_u32le(0)?;
+        }
+        self.bw.write_u32le(0)?; // reserved
+        self.bw.write_u32le(0)?; // reserved
+        self.bw.write_u32le(0)?; // reserved
+        self.bw.write_u32le(0)?; // reserved
+
+        for str in strmgr.iter() {
+            let strl_pos = self.bw.tell() + 8;
+            self.bw.write_buf(b"LIST\0\0\0\0strlstrh")?;
+            self.bw.write_u32le(56)?; // strh size
+
+            match str.get_media_type() {
+                StreamType::Video => {
+                    self.bw.write_buf(b"vids")?;
+                    let fcc = find_avi_fourcc(str.get_info().get_name());
+                    if fcc.is_none() {
+                        return Err(MuxerError::UnsupportedFormat);
+                    }
+                    self.bw.write_buf(&fcc.unwrap_or([0; 4]))?;
+                    let vinfo = str.get_info().get_properties().get_video_info().unwrap();
+                    if vinfo.width >= (1 << 16) || vinfo.height >= (1 << 16) {
+                        return Err(MuxerError::UnsupportedFormat);
+                    }
+                },
+                StreamType::Audio => {
+                    self.bw.write_buf(b"auds")?;
+                    self.bw.write_u32le(0)?;
+                },
+                StreamType::Subtitles => {
+                    self.bw.write_buf(b"txts")?;
+                    self.bw.write_u32le(0)?;
+                },
+                _ => return Err(MuxerError::UnsupportedFormat),
+            };
+            self.stream_info.push(AVIStream {
+                    strh_pos:   self.bw.tell(),
+                    is_video:   str.get_media_type() == StreamType::Video,
+                    nframes:    0,
+                    max_size:   0,
+                });
+
+            self.bw.write_u32le(0)?; // flags
+            self.bw.write_u16le(0)?; // priority
+            self.bw.write_u16le(0)?; // language
+            self.bw.write_u32le(0)?; // initial frames
+            self.bw.write_u32le(str.tb_num)?;
+            self.bw.write_u32le(str.tb_den)?;
+            self.bw.write_u32le(0)?; // start
+            self.bw.write_u32le(0)?; // length
+            self.bw.write_u32le(0)?; // suggested buffer size
+            self.bw.write_u32le(0)?; // quality
+            self.bw.write_u32le(0)?; // sample_size
+            self.bw.write_u16le(0)?; // x
+            self.bw.write_u16le(0)?; // y
+            self.bw.write_u16le(0)?; // w
+            self.bw.write_u16le(0)?; // h
+
+            self.bw.write_buf(b"strf")?;
+            self.bw.write_u32le(0)?;
+            let strf_pos = self.bw.tell();
+            match str.get_media_type() {
+                StreamType::Video => {
+                    let vinfo = str.get_info().get_properties().get_video_info().unwrap();
+                    let hdr_pos = self.bw.tell();
+                    self.bw.write_u32le(0)?;
+                    self.bw.write_u32le(vinfo.width as u32)?;
+                    if vinfo.flipped {
+                        self.bw.write_u32le((-(vinfo.height as i32)) as u32)?;
+                    } else {
+                        self.bw.write_u32le(vinfo.height as u32)?;
+                    }
+                    self.bw.write_u16le(vinfo.format.components as u16)?;
+                    self.bw.write_u16le(vinfo.format.get_total_depth() as u16)?;
+                    let fcc = find_avi_fourcc(str.get_info().get_name());
+                    if fcc.is_none() {
+                        return Err(MuxerError::UnsupportedFormat);
+                    }
+                    self.bw.write_buf(&fcc.unwrap_or([0; 4]))?;
+                    self.bw.write_u32le(0)?; // image size
+                    self.bw.write_u32le(0)?; // x dpi
+                    self.bw.write_u32le(0)?; // y dpi
+                    if vinfo.format.palette {
+//                        unimplemented!();
+                        self.bw.write_u32le(0)?; // total colors
+                        self.bw.write_u32le(0)?; // important colors
+                    } else {
+                        self.bw.write_u32le(0)?; // total colors
+                        self.bw.write_u32le(0)?; // important colors
+                    }
+                    if let Some(ref edata) = str.get_info().get_extradata() {
+                        self.bw.write_buf(edata.as_slice())?;
+                    }
+                    let bisize = self.bw.tell() - hdr_pos;
+                    self.bw.seek(SeekFrom::Current(-(bisize as i64)))?;
+                    self.bw.write_u32le(bisize as u32)?;
+                    self.bw.seek(SeekFrom::End(0))?;
+                },
+                StreamType::Audio => {
+                    let ainfo = str.get_info().get_properties().get_audio_info().unwrap();
+                    let twocc = find_wav_twocc(str.get_info().get_name());
+                    if twocc.is_none() {
+                        return Err(MuxerError::UnsupportedFormat);
+                    }
+                    let twocc = twocc.unwrap_or(0);
+                    self.bw.write_u16le(if twocc == 0 { 1 } else { twocc })?; // PCM hack
+                    self.bw.write_u16le(ainfo.channels as u16)?;
+                    self.bw.write_u32le(ainfo.sample_rate)?;
+                    self.bw.write_u32le(0)?; // avg bytes per second
+                    self.bw.write_u16le(ainfo.block_len as u16)?;
+                    self.bw.write_u16le(ainfo.format.bits as u16)?;
+                    if let Some(ref edata) = str.get_info().get_extradata() {
+                        self.bw.write_buf(edata.as_slice())?;
+                    }
+                },
+                StreamType::Subtitles => {
+                    if let Some(ref edata) = str.get_info().get_extradata() {
+                        self.bw.write_buf(edata.as_slice())?;
+                    }
+                },
+                _ => unreachable!(),
+            };
+            patch_size(&mut self.bw, strf_pos)?;
+            patch_size(&mut self.bw, strl_pos)?;
+        }
+        patch_size(&mut self.bw, hdrl_pos)?;
+
+        self.data_pos = self.bw.tell() + 8;
+        self.bw.write_buf(b"LIST\0\0\0\0movi")?;
+
+        Ok(())
+    }
+    fn mux_frame(&mut self, _strmgr: &StreamManager, pkt: NAPacket) -> MuxerResult<()> {
+        if self.data_pos == 0 {
+            return Err(MuxerError::NotCreated);
+        }
+        let str = pkt.get_stream();
+        let str_num = str.get_num();
+        if str_num > 99 || str_num >= self.stream_info.len() {
+            return Err(MuxerError::UnsupportedFormat);
+        }
+
+        let chunk_len = pkt.get_buffer().len() as u32;
+
+        self.stream_info[str_num].nframes += 1;
+        self.stream_info[str_num].max_size = self.stream_info[str_num].max_size.max(chunk_len);
+// todo palchange
+        self.index.push(IdxEntry {
+                stream: str_num as u32,
+                stype:  str.get_media_type(),
+                key:    pkt.keyframe,
+                pos:    self.bw.tell() as u32,
+                len:    chunk_len });
+        write_chunk_hdr(&mut self.bw, str.get_media_type(), str_num as u32)?;
+        self.bw.write_u32le(chunk_len)?;
+        self.bw.write_buf(pkt.get_buffer().as_slice())?;
+        Ok(())
+    }
+    fn flush(&mut self) -> MuxerResult<()> {
+        Ok(())
+    }
+    fn end(&mut self) -> MuxerResult<()> {
+        patch_size(&mut self.bw, self.data_pos)?;
+        if self.index.len() > 0 {
+            self.bw.write_buf(b"idx1")?;
+            self.bw.write_u32le((self.index.len() * 16) as u32)?;
+            for item in self.index.iter() {
+                write_chunk_hdr(&mut self.bw, item.stype, item.stream)?;
+                self.bw.write_u32le(if item.key { 0x10 } else { 0 })?;
+                self.bw.write_u32le(item.pos)?;
+                self.bw.write_u32le(item.len)?;
+            }
+        }
+        patch_size(&mut self.bw, 8)?;
+        let mut max_frames = 0;
+        let mut max_size = 0;
+        for stri in self.stream_info.iter() {
+            max_frames = max_frames.max(stri.nframes);
+            max_size = max_size.max(stri.max_size);
+            self.bw.seek(SeekFrom::Start(stri.strh_pos + 0x18))?;
+            self.bw.write_u32le(if stri.is_video { stri.nframes } else { 0 })?;
+            self.bw.write_u32le(stri.max_size)?;
+        }
+        self.bw.seek(SeekFrom::Start(0x30))?;
+        self.bw.write_u32le(max_frames)?;
+        self.bw.seek(SeekFrom::Current(8))?;
+        self.bw.write_u32le(max_size)?;
+        Ok(())
+    }
+}
+
+pub struct AVIMuxerCreator {}
+
+impl MuxerCreator for AVIMuxerCreator {
+    fn new_muxer<'a>(&self, bw: &'a mut ByteWriter<'a>) -> Box<dyn MuxCore<'a> + 'a> {
+        Box::new(AVIMuxer::new(bw))
+    }
+    fn get_name(&self) -> &'static str { "avi" }
+}
+
+#[cfg(test)]
+mod test {
+    use super::*;
+    use std::fs::File;
+    use nihav_core::demuxers::*;
+    use crate::demuxers::*;
+
+    #[test]
+    fn test_avi_muxer() {
+        let mut dmx_reg = RegisteredDemuxers::new();
+        generic_register_all_demuxers(&mut dmx_reg);
+        let mut file = File::open("assets/Indeo/laser05.avi").unwrap();
+        let mut fr = FileReader::new_read(&mut file);
+        let mut br = ByteReader::new(&mut fr);
+        let dmx_f = dmx_reg.find_demuxer("avi").unwrap();
+        let mut dmx = create_demuxer(dmx_f, &mut br).unwrap();
+
+        let ofile = File::create("assets/test_out/muxed.avi").unwrap();
+        let mut fw = FileWriter::new_write(ofile);
+        let mut bw = ByteWriter::new(&mut fw);
+        let mut mux = AVIMuxer::new(&mut bw);
+
+        mux.create(dmx.get_stream_manager()).unwrap();
+
+        loop {
+            let pktres = dmx.get_frame();
+            if let Err(e) = pktres {
+                if e == DemuxerError::EOF { break; }
+                panic!("error");
+            }
+            let pkt = pktres.unwrap();
+            println!("Got {}", pkt);
+            mux.mux_frame(dmx.get_stream_manager(), pkt).unwrap();
+        }
+
+        mux.end().unwrap();
+panic!("end");
+    }
+}
diff --git a/nihav-commonfmt/src/muxers/mod.rs b/nihav-commonfmt/src/muxers/mod.rs
new file mode 100644 (file)
index 0000000..6d89806
--- /dev/null
@@ -0,0 +1,15 @@
+use nihav_core::muxers::*;
+
+#[cfg(feature="muxer_avi")]
+mod avi;
+
+const MUXERS: &[&MuxerCreator] = &[
+#[cfg(feature="muxer_avi")]
+    &avi::AVIMuxerCreator {},
+];
+
+pub fn generic_register_all_muxers(rm: &mut RegisteredMuxers) {
+    for muxer in MUXERS.iter() {
+        rm.add_muxer(*muxer);
+    }
+}