| 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 | } |