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     pal_change: bool,
  20 }
  21
  22 struct AVIMuxer<'a> {
  23     bw:             &'a mut ByteWriter<'a>,
  24     index:          Vec<IdxEntry>,
  25     video_str:      Option<usize>,
  26     video_id:       u32,
  27     data_pos:       u64,
  28     stream_info:    Vec<AVIStream>,
  29     pal_pos:        Vec<u32>,
  30 }
  31
  32 impl<'a> AVIMuxer<'a> {
  33     fn new(bw: &'a mut ByteWriter<'a>) -> Self {
  34         Self {
  35             bw,
  36             index:          Vec::new(),
  37             video_str:      None,
  38             video_id:       0,
  39             data_pos:       0,
  40             stream_info:    Vec::with_capacity(2),
  41             pal_pos:        Vec::with_capacity(2),
  42         }
  43     }
  44 }
  45
  46 fn patch_size(bw: &mut ByteWriter, pos: u64) -> MuxerResult<()> {
  47     let size = bw.tell() - pos;
  48     bw.seek(SeekFrom::Current(-((size + 4) as i64)))?;
  49     bw.write_u32le(size as u32)?;
  50     bw.seek(SeekFrom::End(0))?;
  51     Ok(())
  52 }
  53
  54 fn write_chunk_hdr(bw: &mut ByteWriter, stype: StreamType, str_no: u32) -> MuxerResult<()> {
  55     bw.write_byte(b'0' + ((str_no / 10) as u8))?;
  56     bw.write_byte(b'0' + ((str_no % 10) as u8))?;
  57     match stype {
  58         StreamType::Video => { bw.write_buf(b"dc")?; },
  59         StreamType::Audio => { bw.write_buf(b"wb")?; },
  60         StreamType::Subtitles => { bw.write_buf(b"tx")?; },
  61         _ => return Err(MuxerError::UnsupportedFormat),
  62     };
  63     Ok(())
  64 }
  65
  66 impl<'a> MuxCore<'a> for AVIMuxer<'a> {
  67     #[allow(clippy::unreadable_literal)]
  68     #[allow(clippy::cast_lossless)]
  69     fn create(&mut self, strmgr: &StreamManager) -> MuxerResult<()> {
  70         if strmgr.get_num_streams() == 0 {
  71             return Err(MuxerError::InvalidArgument);
  72         }
  73         if strmgr.get_num_streams() > 99 {
  74             return Err(MuxerError::UnsupportedFormat);
  75         }
  76         for (str_no, strm) in strmgr.iter().enumerate() {
  77             if strm.get_media_type() == StreamType::Video {
  78                 self.video_str = Some(str_no);
  79                 self.video_id  = strm.id;
  80                 break;
  81             }
  82         }
  83         let (vinfo, tb_num, tb_den) = if let Some(str_id) = self.video_str {
  84                 let vstr = strmgr.get_stream(str_id).unwrap();
  85                 (vstr.get_info(), vstr.tb_num, vstr.tb_den)
  86             } else {
  87                 (NACodecInfo::new_dummy(), 0, 1)
  88             };
  89         let hdrl_pos = self.bw.tell() + 20;
  90         self.bw.write_buf(b"RIFF\0\0\0\0AVI LIST\0\0\0\0hdrlavih")?;
  91         self.bw.write_u32le(56)?; // avih size
  92         let ms_per_frame = NATimeInfo::ts_to_time(1, 1000000, tb_num, tb_den);
  93         self.bw.write_u32le(ms_per_frame as u32)?;
  94         self.bw.write_u32le(0)?; // max transfer rate
  95         self.bw.write_u32le(0)?; // padding granularity
  96         self.bw.write_u32le(0)?; // flags
  97         self.bw.write_u32le(0)?; // total frames
  98         self.bw.write_u32le(0)?; // initial frames
  99         self.bw.write_u32le(strmgr.get_num_streams() as u32)?;
 100         self.bw.write_u32le(0)?; // suggested buffer size
 101         if let NACodecTypeInfo::Video(ref vinfo) = vinfo.get_properties() {
 102             self.bw.write_u32le(vinfo.width as u32)?;
 103             self.bw.write_u32le(vinfo.height as u32)?;
 104         } else {
 105             self.bw.write_u32le(0)?;
 106             self.bw.write_u32le(0)?;
 107         }
 108         self.bw.write_u32le(0)?; // reserved
 109         self.bw.write_u32le(0)?; // reserved
 110         self.bw.write_u32le(0)?; // reserved
 111         self.bw.write_u32le(0)?; // reserved
 112
 113         self.pal_pos.clear();
 114         self.pal_pos.resize(strmgr.get_num_streams(), 0);
 115         for (strno, strm) in strmgr.iter().enumerate() {
 116             let strl_pos = self.bw.tell() + 8;
 117             self.bw.write_buf(b"LIST\0\0\0\0strlstrh")?;
 118             self.bw.write_u32le(56)?; // strh size
 119
 120             match strm.get_media_type() {
 121                 StreamType::Video => {
 122                     self.bw.write_buf(b"vids")?;
 123                     let fcc = find_avi_fourcc(strm.get_info().get_name());
 124                     if fcc.is_none() {
 125                         return Err(MuxerError::UnsupportedFormat);
 126                     }
 127                     self.bw.write_buf(&fcc.unwrap_or([0; 4]))?;
 128                     let vinfo = strm.get_info().get_properties().get_video_info().unwrap();
 129                     if vinfo.width >= (1 << 16) || vinfo.height >= (1 << 16) {
 130                         return Err(MuxerError::UnsupportedFormat);
 131                     }
 132                 },
 133                 StreamType::Audio => {
 134                     self.bw.write_buf(b"auds")?;
 135                     self.bw.write_u32le(0)?;
 136                 },
 137                 StreamType::Subtitles => {
 138                     self.bw.write_buf(b"txts")?;
 139                     self.bw.write_u32le(0)?;
 140                 },
 141                 _ => return Err(MuxerError::UnsupportedFormat),
 142             };
 143             self.stream_info.push(AVIStream {
 144                     strh_pos:   self.bw.tell(),
 145                     is_video:   strm.get_media_type() == StreamType::Video,
 146                     nframes:    0,
 147                     max_size:   0,
 148                     pal_change: false,
 149                 });
 150
 151             self.bw.write_u32le(0)?; // flags
 152             self.bw.write_u16le(0)?; // priority
 153             self.bw.write_u16le(0)?; // language
 154             self.bw.write_u32le(0)?; // initial frames
 155             self.bw.write_u32le(strm.tb_num)?;
 156             self.bw.write_u32le(strm.tb_den)?;
 157             self.bw.write_u32le(0)?; // start
 158             self.bw.write_u32le(0)?; // length
 159             self.bw.write_u32le(0)?; // suggested buffer size
 160             self.bw.write_u32le(0)?; // quality
 161             self.bw.write_u32le(0)?; // sample_size
 162             self.bw.write_u16le(0)?; // x
 163             self.bw.write_u16le(0)?; // y
 164             self.bw.write_u16le(0)?; // w
 165             self.bw.write_u16le(0)?; // h
 166
 167             self.bw.write_buf(b"strf")?;
 168             self.bw.write_u32le(0)?;
 169             let strf_pos = self.bw.tell();
 170             match strm.get_media_type() {
 171                 StreamType::Video => {
 172                     let vinfo = strm.get_info().get_properties().get_video_info().unwrap();
 173                     let hdr_pos = self.bw.tell();
 174                     self.bw.write_u32le(0)?;
 175                     self.bw.write_u32le(vinfo.width as u32)?;
 176                     self.bw.write_u32le(vinfo.height as u32)?;
 177                     if !vinfo.format.palette {
 178                         self.bw.write_u16le(vinfo.format.components as u16)?;
 179                         self.bw.write_u16le(vinfo.format.get_total_depth() as u16)?;
 180                     } else {
 181                         self.bw.write_u16le(1)?;
 182                         self.bw.write_u16le(8)?;
 183                     }
 184                     let fcc = find_avi_fourcc(strm.get_info().get_name());
 185                     if fcc.is_none() {
 186                         return Err(MuxerError::UnsupportedFormat);
 187                     }
 188                     self.bw.write_buf(&fcc.unwrap_or([0; 4]))?;
 189                     self.bw.write_u32le(0)?; // image size
 190                     self.bw.write_u32le(0)?; // x dpi
 191                     self.bw.write_u32le(0)?; // y dpi
 192                     if vinfo.format.palette {
 193                         self.bw.write_u32le(256)?; // total colors
 194                         self.bw.write_u32le(0)?; // important colors
 195                         self.pal_pos[strno] = self.bw.tell() as u32;
 196                         for _ in 0..256 {
 197                             self.bw.write_u32le(0)?;
 198                         }
 199                     } else {
 200                         self.bw.write_u32le(0)?; // total colors
 201                         self.bw.write_u32le(0)?; // important colors
 202                     }
 203                     if let Some(ref edata) = strm.get_info().get_extradata() {
 204                         self.bw.write_buf(edata.as_slice())?;
 205                     }
 206                     let bisize = self.bw.tell() - hdr_pos;
 207                     self.bw.seek(SeekFrom::Current(-(bisize as i64)))?;
 208                     self.bw.write_u32le(bisize as u32)?;
 209                     self.bw.seek(SeekFrom::End(0))?;
 210                 },
 211                 StreamType::Audio => {
 212                     let ainfo = strm.get_info().get_properties().get_audio_info().unwrap();
 213                     let twocc = find_wav_twocc(strm.get_info().get_name());
 214                     if twocc.is_none() {
 215                         return Err(MuxerError::UnsupportedFormat);
 216                     }
 217                     self.bw.write_u16le(twocc.unwrap_or(0))?;
 218                     self.bw.write_u16le(ainfo.channels as u16)?;
 219                     self.bw.write_u32le(ainfo.sample_rate)?;
 220                     self.bw.write_u32le(0)?; // avg bytes per second
 221                     self.bw.write_u16le(ainfo.block_len as u16)?;
 222                     self.bw.write_u16le(ainfo.format.bits as u16)?;
 223                     if let Some(ref edata) = strm.get_info().get_extradata() {
 224                         self.bw.write_buf(edata.as_slice())?;
 225                     }
 226                 },
 227                 StreamType::Subtitles => {
 228                     if let Some(ref edata) = strm.get_info().get_extradata() {
 229                         self.bw.write_buf(edata.as_slice())?;
 230                     }
 231                 },
 232                 _ => unreachable!(),
 233             };
 234             patch_size(self.bw, strf_pos)?;
 235             patch_size(self.bw, strl_pos)?;
 236         }
 237         patch_size(self.bw, hdrl_pos)?;
 238
 239         self.data_pos = self.bw.tell() + 8;
 240         self.bw.write_buf(b"LIST\0\0\0\0movi")?;
 241
 242         Ok(())
 243     }
 244     fn mux_frame(&mut self, _strmgr: &StreamManager, pkt: NAPacket) -> MuxerResult<()> {
 245         if self.data_pos == 0 {
 246             return Err(MuxerError::NotCreated);
 247         }
 248         let stream = pkt.get_stream();
 249         let str_num = stream.get_num();
 250         if str_num > 99 || str_num >= self.stream_info.len() {
 251             return Err(MuxerError::UnsupportedFormat);
 252         }
 253
 254         let chunk_len = pkt.get_buffer().len() as u32;
 255
 256         if self.pal_pos[str_num] != 0 {
 257             for sdata in pkt.side_data.iter() {
 258                 if let NASideData::Palette(_, ref pal) = sdata {
 259                     let cur_pos = self.bw.tell();
 260                     self.bw.seek(SeekFrom::Start(u64::from(self.pal_pos[str_num])))?;
 261                     for quad in pal.chunks(4) {
 262                         self.bw.write_byte(quad[2])?;
 263                         self.bw.write_byte(quad[1])?;
 264                         self.bw.write_byte(quad[0])?;
 265                         self.bw.write_byte(0)?;
 266                     }
 267                     self.bw.seek(SeekFrom::Start(cur_pos))?;
 268                     self.pal_pos[str_num] = 0;
 269                     break;
 270                 }
 271             }
 272         } else {
 273             for sdata in pkt.side_data.iter() {
 274                 if let NASideData::Palette(true, ref pal) = sdata {
 275                     //todo search for changed region
 276                     let start_clr = 0usize;
 277                     let end_clr = 256usize;
 278                     if start_clr < end_clr {
 279                         let chunk_len = ((end_clr - start_clr) as u32) * 4 + 4;
 280                         self.bw.write_byte(b'0' + ((str_num / 10) as u8))?;
 281                         self.bw.write_byte(b'0' + ((str_num % 10) as u8))?;
 282                         self.bw.write_buf(b"pc")?;
 283                         self.bw.write_u32le(chunk_len)?;
 284                         self.bw.write_byte(start_clr as u8)?;
 285                         self.bw.write_byte((end_clr - start_clr) as u8)?;
 286                         self.bw.write_u16le(0)?; //flags
 287                         self.bw.write_buf(&pal[start_clr * 4..end_clr * 4])?;
 288                         self.stream_info[str_num].pal_change = true;
 289                     }
 290                 }
 291             }
 292         }
 293
 294         self.stream_info[str_num].nframes += 1;
 295         self.stream_info[str_num].max_size = self.stream_info[str_num].max_size.max(chunk_len);
 296         self.index.push(IdxEntry {
 297                 stream: str_num as u32,
 298                 stype:  stream.get_media_type(),
 299                 key:    pkt.keyframe,
 300                 pos:    self.bw.tell() as u32,
 301                 len:    chunk_len });
 302         write_chunk_hdr(self.bw, stream.get_media_type(), str_num as u32)?;
 303         self.bw.write_u32le(chunk_len)?;
 304         self.bw.write_buf(pkt.get_buffer().as_slice())?;
 305         if (self.bw.tell() & 1) != 0 {
 306             self.bw.write_byte(0)?;
 307         }
 308         Ok(())
 309     }
 310     fn flush(&mut self) -> MuxerResult<()> {
 311         Ok(())
 312     }
 313     fn end(&mut self) -> MuxerResult<()> {
 314         patch_size(self.bw, self.data_pos)?;
 315         if !self.index.is_empty() {
 316             self.bw.write_buf(b"idx1")?;
 317             self.bw.write_u32le((self.index.len() * 16) as u32)?;
 318             for item in self.index.iter() {
 319                 write_chunk_hdr(self.bw, item.stype, item.stream)?;
 320                 self.bw.write_u32le(if item.key { 0x10 } else { 0 })?;
 321                 self.bw.write_u32le(item.pos)?;
 322                 self.bw.write_u32le(item.len)?;
 323             }
 324         }
 325         patch_size(self.bw, 8)?;
 326         let mut max_frames = 0;
 327         let mut max_size = 0;
 328         for stri in self.stream_info.iter() {
 329             max_frames = max_frames.max(stri.nframes);
 330             max_size = max_size.max(stri.max_size);
 331             if stri.pal_change {
 332                 self.bw.seek(SeekFrom::Start(stri.strh_pos))?;
 333                 self.bw.write_u32le(0x00010000)?;
 334             }
 335             self.bw.seek(SeekFrom::Start(stri.strh_pos + 0x18))?;
 336             self.bw.write_u32le(if stri.is_video { stri.nframes } else { 0 })?;
 337             self.bw.write_u32le(stri.max_size)?;
 338         }
 339         self.bw.seek(SeekFrom::Start(0x30))?;
 340         self.bw.write_u32le(max_frames)?;
 341         self.bw.seek(SeekFrom::Current(8))?;
 342         self.bw.write_u32le(max_size)?;
 343         Ok(())
 344     }
 345 }
 346
 347 impl<'a> NAOptionHandler for AVIMuxer<'a> {
 348     fn get_supported_options(&self) -> &[NAOptionDefinition] { &[] }
 349     fn set_options(&mut self, _options: &[NAOption]) { }
 350     fn query_option_value(&self, _name: &str) -> Option<NAValue> { None }
 351 }
 352
 353 pub struct AVIMuxerCreator {}
 354
 355 impl MuxerCreator for AVIMuxerCreator {
 356     fn new_muxer<'a>(&self, bw: &'a mut ByteWriter<'a>) -> Box<dyn MuxCore<'a> + 'a> {
 357         Box::new(AVIMuxer::new(bw))
 358     }
 359     fn get_name(&self) -> &'static str { "avi" }
 360     fn get_capabilities(&self) -> MuxerCapabilities { MuxerCapabilities::Universal }
 361 }
 362
 363 #[cfg(test)]
 364 mod test {
 365     use nihav_core::codecs::*;
 366     use nihav_core::demuxers::*;
 367     use nihav_core::muxers::*;
 368     use nihav_codec_support::test::enc_video::*;
 369     use crate::*;
 370
 371     #[test]
 372     fn test_avi_muxer() {
 373         let mut dmx_reg = RegisteredDemuxers::new();
 374         generic_register_all_demuxers(&mut dmx_reg);
 375         //test sample: https://samples.mplayerhq.hu/V-codecs/RT21/320x240/laser05.avi
 376         let dec_config = DecoderTestParams {
 377                 demuxer:        "avi",
 378                 in_name:        "assets/Indeo/laser05.avi",
 379                 limit:          None,
 380                 stream_type:    StreamType::None,
 381                 dmx_reg, dec_reg: RegisteredDecoders::new(),
 382             };
 383         let mut mux_reg = RegisteredMuxers::new();
 384         generic_register_all_muxers(&mut mux_reg);
 385         /*let enc_config = EncoderTestParams {
 386                 muxer:      "avi",
 387                 enc_name:   "",
 388                 out_name:   "muxed.avi",
 389                 mux_reg, enc_reg: RegisteredEncoders::new(),
 390             };
 391         test_remuxing(&dec_config, &enc_config);*/
 392         test_remuxing_md5(&dec_config, "avi", &mux_reg,
 393                           [0xa0fb0e47, 0x412e24dd, 0x6b89711c, 0x276fb799]);
 394     }
 395 }