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