nihav-commonfmt/src/muxers/avi.rs

   1 use nihav_core::muxers::*;
   2 use nihav_registry::register::*;
   3
   4 #[derive(Clone,Copy)]
   5 struct IdxEntry {
   6     stream:     u32,
   7     stype:      StreamType,
   8     key:        bool,
   9     pos:        u32,
  10     len:        u32,
  11 }
  12
  13 #[derive(Clone,Copy)]
  14 struct AVIStream {
  15     strh_pos:   u64,
  16     nframes:    u32,
  17     is_video:   bool,
  18     max_size:   u32,
  19 }
  20
  21 struct AVIMuxer<'a> {
  22     bw:             &'a mut ByteWriter<'a>,
  23     index:          Vec<IdxEntry>,
  24     video_str:      Option<usize>,
  25     video_id:       u32,
  26     data_pos:       u64,
  27     stream_info:    Vec<AVIStream>,
  28 }
  29
  30 impl<'a> AVIMuxer<'a> {
  31     fn new(bw: &'a mut ByteWriter<'a>) -> Self {
  32         Self {
  33             bw,
  34             index:          Vec::new(),
  35             video_str:      None,
  36             video_id:       0,
  37             data_pos:       0,
  38             stream_info:    Vec::with_capacity(2),
  39         }
  40     }
  41 }
  42
  43 fn patch_size(bw: &mut ByteWriter, pos: u64) -> MuxerResult<()> {
  44     let size = bw.tell() - pos;
  45     bw.seek(SeekFrom::Current(-((size + 4) as i64)))?;
  46     bw.write_u32le(size as u32)?;
  47     bw.seek(SeekFrom::End(0))?;
  48     Ok(())
  49 }
  50
  51 fn write_chunk_hdr(bw: &mut ByteWriter, stype: StreamType, str_no: u32) -> MuxerResult<()> {
  52     bw.write_byte(b'0' + ((str_no / 10) as u8))?;
  53     bw.write_byte(b'0' + ((str_no % 10) as u8))?;
  54     match stype {
  55         StreamType::Video => { bw.write_buf(b"dc")?; },
  56         StreamType::Audio => { bw.write_buf(b"wb")?; },
  57         StreamType::Subtitles => { bw.write_buf(b"tx")?; },
  58         _ => return Err(MuxerError::UnsupportedFormat),
  59     };
  60     Ok(())
  61 }
  62
  63 impl<'a> MuxCore<'a> for AVIMuxer<'a> {
  64     fn create(&mut self, strmgr: &StreamManager) -> MuxerResult<()> {
  65         if strmgr.get_num_streams() == 0 {
  66             return Err(MuxerError::InvalidArgument);
  67         }
  68         if strmgr.get_num_streams() > 99 {
  69             return Err(MuxerError::UnsupportedFormat);
  70         }
  71         for (str_no, str) in strmgr.iter().enumerate() {
  72             if str.get_media_type() == StreamType::Video {
  73                 self.video_str = Some(str_no);
  74                 self.video_id  = str.id;
  75                 break;
  76             }
  77         }
  78         let (vinfo, tb_num, tb_den) = if let Some(str_id) = self.video_str {
  79                 let vstr = strmgr.get_stream(str_id).unwrap();
  80                 (vstr.get_info(), vstr.tb_num, vstr.tb_den)
  81             } else {
  82                 (NACodecInfo::new_dummy(), 0, 1)
  83             };
  84         let hdrl_pos = self.bw.tell() + 20;
  85         self.bw.write_buf(b"RIFF\0\0\0\0AVI LIST\0\0\0\0hdrlavih")?;
  86         self.bw.write_u32le(56)?; // avih size
  87         let ms_per_frame = NATimeInfo::ts_to_time(1, 1000000, tb_num, tb_den);
  88         self.bw.write_u32le(ms_per_frame as u32)?;
  89         self.bw.write_u32le(0)?; // max transfer rate
  90         self.bw.write_u32le(0)?; // padding granularity
  91         self.bw.write_u32le(0)?; // flags
  92         self.bw.write_u32le(0)?; // total frames
  93         self.bw.write_u32le(0)?; // initial frames
  94         self.bw.write_u32le(strmgr.get_num_streams() as u32)?;
  95         self.bw.write_u32le(0)?; // suggested buffer size
  96         if let NACodecTypeInfo::Video(ref vinfo) = vinfo.get_properties() {
  97             self.bw.write_u32le(vinfo.width as u32)?;
  98             self.bw.write_u32le(vinfo.height as u32)?;
  99         } else {
 100             self.bw.write_u32le(0)?;
 101             self.bw.write_u32le(0)?;
 102         }
 103         self.bw.write_u32le(0)?; // reserved
 104         self.bw.write_u32le(0)?; // reserved
 105         self.bw.write_u32le(0)?; // reserved
 106         self.bw.write_u32le(0)?; // reserved
 107
 108         for str in strmgr.iter() {
 109             let strl_pos = self.bw.tell() + 8;
 110             self.bw.write_buf(b"LIST\0\0\0\0strlstrh")?;
 111             self.bw.write_u32le(56)?; // strh size
 112
 113             match str.get_media_type() {
 114                 StreamType::Video => {
 115                     self.bw.write_buf(b"vids")?;
 116                     let fcc = find_avi_fourcc(str.get_info().get_name());
 117                     if fcc.is_none() {
 118                         return Err(MuxerError::UnsupportedFormat);
 119                     }
 120                     self.bw.write_buf(&fcc.unwrap_or([0; 4]))?;
 121                     let vinfo = str.get_info().get_properties().get_video_info().unwrap();
 122                     if vinfo.width >= (1 << 16) || vinfo.height >= (1 << 16) {
 123                         return Err(MuxerError::UnsupportedFormat);
 124                     }
 125                 },
 126                 StreamType::Audio => {
 127                     self.bw.write_buf(b"auds")?;
 128                     self.bw.write_u32le(0)?;
 129                 },
 130                 StreamType::Subtitles => {
 131                     self.bw.write_buf(b"txts")?;
 132                     self.bw.write_u32le(0)?;
 133                 },
 134                 _ => return Err(MuxerError::UnsupportedFormat),
 135             };
 136             self.stream_info.push(AVIStream {
 137                     strh_pos:   self.bw.tell(),
 138                     is_video:   str.get_media_type() == StreamType::Video,
 139                     nframes:    0,
 140                     max_size:   0,
 141                 });
 142
 143             self.bw.write_u32le(0)?; // flags
 144             self.bw.write_u16le(0)?; // priority
 145             self.bw.write_u16le(0)?; // language
 146             self.bw.write_u32le(0)?; // initial frames
 147             self.bw.write_u32le(str.tb_num)?;
 148             self.bw.write_u32le(str.tb_den)?;
 149             self.bw.write_u32le(0)?; // start
 150             self.bw.write_u32le(0)?; // length
 151             self.bw.write_u32le(0)?; // suggested buffer size
 152             self.bw.write_u32le(0)?; // quality
 153             self.bw.write_u32le(0)?; // sample_size
 154             self.bw.write_u16le(0)?; // x
 155             self.bw.write_u16le(0)?; // y
 156             self.bw.write_u16le(0)?; // w
 157             self.bw.write_u16le(0)?; // h
 158
 159             self.bw.write_buf(b"strf")?;
 160             self.bw.write_u32le(0)?;
 161             let strf_pos = self.bw.tell();
 162             match str.get_media_type() {
 163                 StreamType::Video => {
 164                     let vinfo = str.get_info().get_properties().get_video_info().unwrap();
 165                     let hdr_pos = self.bw.tell();
 166                     self.bw.write_u32le(0)?;
 167                     self.bw.write_u32le(vinfo.width as u32)?;
 168                     if vinfo.flipped {
 169                         self.bw.write_u32le((-(vinfo.height as i32)) as u32)?;
 170                     } else {
 171                         self.bw.write_u32le(vinfo.height as u32)?;
 172                     }
 173                     self.bw.write_u16le(vinfo.format.components as u16)?;
 174                     self.bw.write_u16le(vinfo.format.get_total_depth() as u16)?;
 175                     let fcc = find_avi_fourcc(str.get_info().get_name());
 176                     if fcc.is_none() {
 177                         return Err(MuxerError::UnsupportedFormat);
 178                     }
 179                     self.bw.write_buf(&fcc.unwrap_or([0; 4]))?;
 180                     self.bw.write_u32le(0)?; // image size
 181                     self.bw.write_u32le(0)?; // x dpi
 182                     self.bw.write_u32le(0)?; // y dpi
 183                     if vinfo.format.palette {
 184 //                        unimplemented!();
 185                         self.bw.write_u32le(0)?; // total colors
 186                         self.bw.write_u32le(0)?; // important colors
 187                     } else {
 188                         self.bw.write_u32le(0)?; // total colors
 189                         self.bw.write_u32le(0)?; // important colors
 190                     }
 191                     if let Some(ref edata) = str.get_info().get_extradata() {
 192                         self.bw.write_buf(edata.as_slice())?;
 193                     }
 194                     let bisize = self.bw.tell() - hdr_pos;
 195                     self.bw.seek(SeekFrom::Current(-(bisize as i64)))?;
 196                     self.bw.write_u32le(bisize as u32)?;
 197                     self.bw.seek(SeekFrom::End(0))?;
 198                 },
 199                 StreamType::Audio => {
 200                     let ainfo = str.get_info().get_properties().get_audio_info().unwrap();
 201                     let twocc = find_wav_twocc(str.get_info().get_name());
 202                     if twocc.is_none() {
 203                         return Err(MuxerError::UnsupportedFormat);
 204                     }
 205                     self.bw.write_u16le(twocc.unwrap_or(0))?;
 206                     self.bw.write_u16le(ainfo.channels as u16)?;
 207                     self.bw.write_u32le(ainfo.sample_rate)?;
 208                     self.bw.write_u32le(0)?; // avg bytes per second
 209                     self.bw.write_u16le(ainfo.block_len as u16)?;
 210                     self.bw.write_u16le(ainfo.format.bits as u16)?;
 211                     if let Some(ref edata) = str.get_info().get_extradata() {
 212                         self.bw.write_buf(edata.as_slice())?;
 213                     }
 214                 },
 215                 StreamType::Subtitles => {
 216                     if let Some(ref edata) = str.get_info().get_extradata() {
 217                         self.bw.write_buf(edata.as_slice())?;
 218                     }
 219                 },
 220                 _ => unreachable!(),
 221             };
 222             patch_size(&mut self.bw, strf_pos)?;
 223             patch_size(&mut self.bw, strl_pos)?;
 224         }
 225         patch_size(&mut self.bw, hdrl_pos)?;
 226
 227         self.data_pos = self.bw.tell() + 8;
 228         self.bw.write_buf(b"LIST\0\0\0\0movi")?;
 229
 230         Ok(())
 231     }
 232     fn mux_frame(&mut self, _strmgr: &StreamManager, pkt: NAPacket) -> MuxerResult<()> {
 233         if self.data_pos == 0 {
 234             return Err(MuxerError::NotCreated);
 235         }
 236         let str = pkt.get_stream();
 237         let str_num = str.get_num();
 238         if str_num > 99 || str_num >= self.stream_info.len() {
 239             return Err(MuxerError::UnsupportedFormat);
 240         }
 241
 242         let chunk_len = pkt.get_buffer().len() as u32;
 243
 244         self.stream_info[str_num].nframes += 1;
 245         self.stream_info[str_num].max_size = self.stream_info[str_num].max_size.max(chunk_len);
 246 // todo palchange
 247         self.index.push(IdxEntry {
 248                 stream: str_num as u32,
 249                 stype:  str.get_media_type(),
 250                 key:    pkt.keyframe,
 251                 pos:    self.bw.tell() as u32,
 252                 len:    chunk_len });
 253         write_chunk_hdr(&mut self.bw, str.get_media_type(), str_num as u32)?;
 254         self.bw.write_u32le(chunk_len)?;
 255         self.bw.write_buf(pkt.get_buffer().as_slice())?;
 256         Ok(())
 257     }
 258     fn flush(&mut self) -> MuxerResult<()> {
 259         Ok(())
 260     }
 261     fn end(&mut self) -> MuxerResult<()> {
 262         patch_size(&mut self.bw, self.data_pos)?;
 263         if self.index.len() > 0 {
 264             self.bw.write_buf(b"idx1")?;
 265             self.bw.write_u32le((self.index.len() * 16) as u32)?;
 266             for item in self.index.iter() {
 267                 write_chunk_hdr(&mut self.bw, item.stype, item.stream)?;
 268                 self.bw.write_u32le(if item.key { 0x10 } else { 0 })?;
 269                 self.bw.write_u32le(item.pos)?;
 270                 self.bw.write_u32le(item.len)?;
 271             }
 272         }
 273         patch_size(&mut self.bw, 8)?;
 274         let mut max_frames = 0;
 275         let mut max_size = 0;
 276         for stri in self.stream_info.iter() {
 277             max_frames = max_frames.max(stri.nframes);
 278             max_size = max_size.max(stri.max_size);
 279             self.bw.seek(SeekFrom::Start(stri.strh_pos + 0x18))?;
 280             self.bw.write_u32le(if stri.is_video { stri.nframes } else { 0 })?;
 281             self.bw.write_u32le(stri.max_size)?;
 282         }
 283         self.bw.seek(SeekFrom::Start(0x30))?;
 284         self.bw.write_u32le(max_frames)?;
 285         self.bw.seek(SeekFrom::Current(8))?;
 286         self.bw.write_u32le(max_size)?;
 287         Ok(())
 288     }
 289 }
 290
 291 impl<'a> NAOptionHandler for AVIMuxer<'a> {
 292     fn get_supported_options(&self) -> &[NAOptionDefinition] { &[] }
 293     fn set_options(&mut self, _options: &[NAOption]) { }
 294     fn query_option_value(&self, _name: &str) -> Option<NAValue> { None }
 295 }
 296
 297 pub struct AVIMuxerCreator {}
 298
 299 impl MuxerCreator for AVIMuxerCreator {
 300     fn new_muxer<'a>(&self, bw: &'a mut ByteWriter<'a>) -> Box<dyn MuxCore<'a> + 'a> {
 301         Box::new(AVIMuxer::new(bw))
 302     }
 303     fn get_name(&self) -> &'static str { "avi" }
 304     fn get_capabilities(&self) -> MuxerCapabilities { MuxerCapabilities::Universal }
 305 }
 306
 307 #[cfg(test)]
 308 mod test {
 309     use super::*;
 310     use std::fs::File;
 311     use nihav_core::demuxers::*;
 312     use crate::demuxers::*;
 313
 314     #[test]
 315     fn test_avi_muxer() {
 316         let mut dmx_reg = RegisteredDemuxers::new();
 317         generic_register_all_demuxers(&mut dmx_reg);
 318         let mut file = File::open("assets/Indeo/laser05.avi").unwrap();
 319         let mut fr = FileReader::new_read(&mut file);
 320         let mut br = ByteReader::new(&mut fr);
 321         let dmx_f = dmx_reg.find_demuxer("avi").unwrap();
 322         let mut dmx = create_demuxer(dmx_f, &mut br).unwrap();
 323
 324         let ofile = File::create("assets/test_out/muxed.avi").unwrap();
 325         let mut fw = FileWriter::new_write(ofile);
 326         let mut bw = ByteWriter::new(&mut fw);
 327         let mut mux = AVIMuxer::new(&mut bw);
 328
 329         mux.create(dmx.get_stream_manager()).unwrap();
 330
 331         loop {
 332             let pktres = dmx.get_frame();
 333             if let Err(e) = pktres {
 334                 if e == DemuxerError::EOF { break; }
 335                 panic!("error");
 336             }
 337             let pkt = pktres.unwrap();
 338             println!("Got {}", pkt);
 339             mux.mux_frame(dmx.get_stream_manager(), pkt).unwrap();
 340         }
 341
 342         mux.end().unwrap();
 343 panic!("end");
 344     }
 345 }