]>
Commit | Line | Data |
---|---|---|
5cd2f537 KS |
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, | |
c8db9313 | 18 | max_size: u32, |
5cd2f537 KS |
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> { | |
61cab15b KS |
64 | #[allow(clippy::unreadable_literal)] |
65 | #[allow(clippy::cast_lossless)] | |
5cd2f537 KS |
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 | |
61cab15b | 189 | println!("pal?"); |
5cd2f537 KS |
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 | } | |
586ea187 | 208 | self.bw.write_u16le(twocc.unwrap_or(0))?; |
5cd2f537 KS |
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())?; | |
77d52ff6 KS |
259 | if (self.bw.tell() & 1) != 0 { |
260 | self.bw.write_byte(0)?; | |
261 | } | |
5cd2f537 KS |
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)?; | |
61cab15b | 269 | if !self.index.is_empty() { |
5cd2f537 KS |
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 | ||
dc80f48e KS |
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 | ||
5cd2f537 KS |
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" } | |
f0081142 | 310 | fn get_capabilities(&self) -> MuxerCapabilities { MuxerCapabilities::Universal } |
5cd2f537 KS |
311 | } |
312 | ||
313 | #[cfg(test)] | |
314 | mod test { | |
68e5a4ca | 315 | use nihav_core::codecs::*; |
5cd2f537 | 316 | use nihav_core::demuxers::*; |
68e5a4ca KS |
317 | use nihav_core::muxers::*; |
318 | use nihav_codec_support::test::enc_video::*; | |
319 | use crate::*; | |
5cd2f537 KS |
320 | |
321 | #[test] | |
322 | fn test_avi_muxer() { | |
323 | let mut dmx_reg = RegisteredDemuxers::new(); | |
324 | generic_register_all_demuxers(&mut dmx_reg); | |
68e5a4ca KS |
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]); | |
5cd2f537 KS |
343 | } |
344 | } |