From 5cd2f53775f5dcb021ab1294417a57d4fc303d42 Mon Sep 17 00:00:00 2001 From: Kostya Shishkov Date: Thu, 21 May 2020 11:01:00 +0200 Subject: [PATCH] avi muxer --- nihav-commonfmt/Cargo.toml | 5 +- nihav-commonfmt/src/lib.rs | 5 + nihav-commonfmt/src/muxers/avi.rs | 339 ++++++++++++++++++++++++++++++ nihav-commonfmt/src/muxers/mod.rs | 15 ++ 4 files changed, 363 insertions(+), 1 deletion(-) create mode 100644 nihav-commonfmt/src/muxers/avi.rs create mode 100644 nihav-commonfmt/src/muxers/mod.rs diff --git a/nihav-commonfmt/Cargo.toml b/nihav-commonfmt/Cargo.toml index 6c0e810..5807da8 100644 --- a/nihav-commonfmt/Cargo.toml +++ b/nihav-commonfmt/Cargo.toml @@ -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"] diff --git a/nihav-commonfmt/src/lib.rs b/nihav-commonfmt/src/lib.rs index ef5ad07..82bfd49 100644 --- a/nihav-commonfmt/src/lib.rs +++ b/nihav-commonfmt/src/lib.rs @@ -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 index 0000000..d47aac3 --- /dev/null +++ b/nihav-commonfmt/src/muxers/avi.rs @@ -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, + video_str: Option, + video_id: u32, + data_pos: u64, + stream_info: Vec, +} + +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 + '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 index 0000000..6d89806 --- /dev/null +++ b/nihav-commonfmt/src/muxers/mod.rs @@ -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); + } +} -- 2.39.5