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