]>
Commit | Line | Data |
---|---|---|
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, strm) in strmgr.iter().enumerate() { | |
76 | if strm.get_media_type() == StreamType::Video { | |
77 | self.video_str = Some(str_no); | |
78 | self.video_id = strm.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, strm) 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 strm.get_media_type() { | |
120 | StreamType::Video => { | |
121 | self.bw.write_buf(b"vids")?; | |
122 | let fcc = find_avi_fourcc(strm.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 = strm.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: strm.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(strm.tb_num)?; | |
154 | self.bw.write_u32le(strm.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 strm.get_media_type() { | |
169 | StreamType::Video => { | |
170 | let vinfo = strm.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 | if !vinfo.format.palette { | |
176 | self.bw.write_u16le(vinfo.format.components as u16)?; | |
177 | self.bw.write_u16le(vinfo.format.get_total_depth() as u16)?; | |
178 | } else { | |
179 | self.bw.write_u16le(1)?; | |
180 | self.bw.write_u16le(8)?; | |
181 | } | |
182 | let fcc = find_avi_fourcc(strm.get_info().get_name()); | |
183 | if fcc.is_none() { | |
184 | return Err(MuxerError::UnsupportedFormat); | |
185 | } | |
186 | self.bw.write_buf(&fcc.unwrap_or([0; 4]))?; | |
187 | self.bw.write_u32le(0)?; // image size | |
188 | self.bw.write_u32le(0)?; // x dpi | |
189 | self.bw.write_u32le(0)?; // y dpi | |
190 | if vinfo.format.palette { | |
191 | self.bw.write_u32le(256)?; // total colors | |
192 | self.bw.write_u32le(0)?; // important colors | |
193 | self.pal_pos[strno] = self.bw.tell() as u32; | |
194 | for _ in 0..256 { | |
195 | self.bw.write_u32le(0)?; | |
196 | } | |
197 | } else { | |
198 | self.bw.write_u32le(0)?; // total colors | |
199 | self.bw.write_u32le(0)?; // important colors | |
200 | } | |
201 | if let Some(ref edata) = strm.get_info().get_extradata() { | |
202 | self.bw.write_buf(edata.as_slice())?; | |
203 | } | |
204 | let bisize = self.bw.tell() - hdr_pos; | |
205 | self.bw.seek(SeekFrom::Current(-(bisize as i64)))?; | |
206 | self.bw.write_u32le(bisize as u32)?; | |
207 | self.bw.seek(SeekFrom::End(0))?; | |
208 | }, | |
209 | StreamType::Audio => { | |
210 | let ainfo = strm.get_info().get_properties().get_audio_info().unwrap(); | |
211 | let twocc = find_wav_twocc(strm.get_info().get_name()); | |
212 | if twocc.is_none() { | |
213 | return Err(MuxerError::UnsupportedFormat); | |
214 | } | |
215 | self.bw.write_u16le(twocc.unwrap_or(0))?; | |
216 | self.bw.write_u16le(ainfo.channels as u16)?; | |
217 | self.bw.write_u32le(ainfo.sample_rate)?; | |
218 | self.bw.write_u32le(0)?; // avg bytes per second | |
219 | self.bw.write_u16le(ainfo.block_len as u16)?; | |
220 | self.bw.write_u16le(ainfo.format.bits as u16)?; | |
221 | if let Some(ref edata) = strm.get_info().get_extradata() { | |
222 | self.bw.write_buf(edata.as_slice())?; | |
223 | } | |
224 | }, | |
225 | StreamType::Subtitles => { | |
226 | if let Some(ref edata) = strm.get_info().get_extradata() { | |
227 | self.bw.write_buf(edata.as_slice())?; | |
228 | } | |
229 | }, | |
230 | _ => unreachable!(), | |
231 | }; | |
232 | patch_size(self.bw, strf_pos)?; | |
233 | patch_size(self.bw, strl_pos)?; | |
234 | } | |
235 | patch_size(self.bw, hdrl_pos)?; | |
236 | ||
237 | self.data_pos = self.bw.tell() + 8; | |
238 | self.bw.write_buf(b"LIST\0\0\0\0movi")?; | |
239 | ||
240 | Ok(()) | |
241 | } | |
242 | fn mux_frame(&mut self, _strmgr: &StreamManager, pkt: NAPacket) -> MuxerResult<()> { | |
243 | if self.data_pos == 0 { | |
244 | return Err(MuxerError::NotCreated); | |
245 | } | |
246 | let stream = pkt.get_stream(); | |
247 | let str_num = stream.get_num(); | |
248 | if str_num > 99 || str_num >= self.stream_info.len() { | |
249 | return Err(MuxerError::UnsupportedFormat); | |
250 | } | |
251 | ||
252 | let chunk_len = pkt.get_buffer().len() as u32; | |
253 | ||
254 | if self.pal_pos[str_num] != 0 { | |
255 | for sdata in pkt.side_data.iter() { | |
256 | if let NASideData::Palette(_, ref pal) = sdata { | |
257 | let cur_pos = self.bw.tell(); | |
258 | self.bw.seek(SeekFrom::Start(u64::from(self.pal_pos[str_num])))?; | |
259 | self.bw.write_buf(pal.as_ref())?; | |
260 | self.bw.seek(SeekFrom::Start(cur_pos))?; | |
261 | self.pal_pos[str_num] = 0; | |
262 | break; | |
263 | } | |
264 | } | |
265 | } else { | |
266 | for sdata in pkt.side_data.iter() { | |
267 | if let NASideData::Palette(true, ref pal) = sdata { | |
268 | //todo search for changed region | |
269 | let start_clr = 0usize; | |
270 | let end_clr = 256usize; | |
271 | if start_clr < end_clr { | |
272 | let chunk_len = ((end_clr - start_clr) as u32) * 4 + 4; | |
273 | self.bw.write_byte(b'0' + ((str_num / 10) as u8))?; | |
274 | self.bw.write_byte(b'0' + ((str_num % 10) as u8))?; | |
275 | self.bw.write_buf(b"pc")?; | |
276 | self.bw.write_u32le(chunk_len)?; | |
277 | self.bw.write_byte(start_clr as u8)?; | |
278 | self.bw.write_byte((end_clr - start_clr) as u8)?; | |
279 | self.bw.write_u16le(0)?; //flags | |
280 | self.bw.write_buf(&pal[start_clr * 4..end_clr * 4])?; | |
281 | } | |
282 | } | |
283 | } | |
284 | } | |
285 | ||
286 | self.stream_info[str_num].nframes += 1; | |
287 | self.stream_info[str_num].max_size = self.stream_info[str_num].max_size.max(chunk_len); | |
288 | self.index.push(IdxEntry { | |
289 | stream: str_num as u32, | |
290 | stype: stream.get_media_type(), | |
291 | key: pkt.keyframe, | |
292 | pos: self.bw.tell() as u32, | |
293 | len: chunk_len }); | |
294 | write_chunk_hdr(self.bw, stream.get_media_type(), str_num as u32)?; | |
295 | self.bw.write_u32le(chunk_len)?; | |
296 | self.bw.write_buf(pkt.get_buffer().as_slice())?; | |
297 | if (self.bw.tell() & 1) != 0 { | |
298 | self.bw.write_byte(0)?; | |
299 | } | |
300 | Ok(()) | |
301 | } | |
302 | fn flush(&mut self) -> MuxerResult<()> { | |
303 | Ok(()) | |
304 | } | |
305 | fn end(&mut self) -> MuxerResult<()> { | |
306 | patch_size(self.bw, self.data_pos)?; | |
307 | if !self.index.is_empty() { | |
308 | self.bw.write_buf(b"idx1")?; | |
309 | self.bw.write_u32le((self.index.len() * 16) as u32)?; | |
310 | for item in self.index.iter() { | |
311 | write_chunk_hdr(self.bw, item.stype, item.stream)?; | |
312 | self.bw.write_u32le(if item.key { 0x10 } else { 0 })?; | |
313 | self.bw.write_u32le(item.pos)?; | |
314 | self.bw.write_u32le(item.len)?; | |
315 | } | |
316 | } | |
317 | patch_size(self.bw, 8)?; | |
318 | let mut max_frames = 0; | |
319 | let mut max_size = 0; | |
320 | for stri in self.stream_info.iter() { | |
321 | max_frames = max_frames.max(stri.nframes); | |
322 | max_size = max_size.max(stri.max_size); | |
323 | self.bw.seek(SeekFrom::Start(stri.strh_pos + 0x18))?; | |
324 | self.bw.write_u32le(if stri.is_video { stri.nframes } else { 0 })?; | |
325 | self.bw.write_u32le(stri.max_size)?; | |
326 | } | |
327 | self.bw.seek(SeekFrom::Start(0x30))?; | |
328 | self.bw.write_u32le(max_frames)?; | |
329 | self.bw.seek(SeekFrom::Current(8))?; | |
330 | self.bw.write_u32le(max_size)?; | |
331 | Ok(()) | |
332 | } | |
333 | } | |
334 | ||
335 | impl<'a> NAOptionHandler for AVIMuxer<'a> { | |
336 | fn get_supported_options(&self) -> &[NAOptionDefinition] { &[] } | |
337 | fn set_options(&mut self, _options: &[NAOption]) { } | |
338 | fn query_option_value(&self, _name: &str) -> Option<NAValue> { None } | |
339 | } | |
340 | ||
341 | pub struct AVIMuxerCreator {} | |
342 | ||
343 | impl MuxerCreator for AVIMuxerCreator { | |
344 | fn new_muxer<'a>(&self, bw: &'a mut ByteWriter<'a>) -> Box<dyn MuxCore<'a> + 'a> { | |
345 | Box::new(AVIMuxer::new(bw)) | |
346 | } | |
347 | fn get_name(&self) -> &'static str { "avi" } | |
348 | fn get_capabilities(&self) -> MuxerCapabilities { MuxerCapabilities::Universal } | |
349 | } | |
350 | ||
351 | #[cfg(test)] | |
352 | mod test { | |
353 | use nihav_core::codecs::*; | |
354 | use nihav_core::demuxers::*; | |
355 | use nihav_core::muxers::*; | |
356 | use nihav_codec_support::test::enc_video::*; | |
357 | use crate::*; | |
358 | ||
359 | #[test] | |
360 | fn test_avi_muxer() { | |
361 | let mut dmx_reg = RegisteredDemuxers::new(); | |
362 | generic_register_all_demuxers(&mut dmx_reg); | |
363 | //test sample: https://samples.mplayerhq.hu/V-codecs/RT21/320x240/laser05.avi | |
364 | let dec_config = DecoderTestParams { | |
365 | demuxer: "avi", | |
366 | in_name: "assets/Indeo/laser05.avi", | |
367 | limit: None, | |
368 | stream_type: StreamType::None, | |
369 | dmx_reg, dec_reg: RegisteredDecoders::new(), | |
370 | }; | |
371 | let mut mux_reg = RegisteredMuxers::new(); | |
372 | generic_register_all_muxers(&mut mux_reg); | |
373 | /*let enc_config = EncoderTestParams { | |
374 | muxer: "avi", | |
375 | enc_name: "", | |
376 | out_name: "muxed.avi", | |
377 | mux_reg, enc_reg: RegisteredEncoders::new(), | |
378 | }; | |
379 | test_remuxing(&dec_config, &enc_config);*/ | |
380 | test_remuxing_md5(&dec_config, "avi", &mux_reg, | |
381 | [0xa0fb0e47, 0x412e24dd, 0x6b89711c, 0x276fb799]); | |
382 | } | |
383 | } |