]>
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> { | |
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 | } | |
586ea187 | 205 | self.bw.write_u16le(twocc.unwrap_or(0))?; |
5cd2f537 KS |
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 | ||
dc80f48e KS |
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 | ||
5cd2f537 KS |
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" } | |
f0081142 | 304 | fn get_capabilities(&self) -> MuxerCapabilities { MuxerCapabilities::Universal } |
5cd2f537 KS |
305 | } |
306 | ||
307 | #[cfg(test)] | |
308 | mod test { | |
68e5a4ca | 309 | use nihav_core::codecs::*; |
5cd2f537 | 310 | use nihav_core::demuxers::*; |
68e5a4ca KS |
311 | use nihav_core::muxers::*; |
312 | use nihav_codec_support::test::enc_video::*; | |
313 | use crate::*; | |
5cd2f537 KS |
314 | |
315 | #[test] | |
316 | fn test_avi_muxer() { | |
317 | let mut dmx_reg = RegisteredDemuxers::new(); | |
318 | generic_register_all_demuxers(&mut dmx_reg); | |
68e5a4ca KS |
319 | let dec_config = DecoderTestParams { |
320 | demuxer: "avi", | |
321 | in_name: "assets/Indeo/laser05.avi", | |
322 | limit: None, | |
323 | stream_type: StreamType::None, | |
324 | dmx_reg, dec_reg: RegisteredDecoders::new(), | |
325 | }; | |
326 | let mut mux_reg = RegisteredMuxers::new(); | |
327 | generic_register_all_muxers(&mut mux_reg); | |
328 | /*let enc_config = EncoderTestParams { | |
329 | muxer: "avi", | |
330 | enc_name: "", | |
331 | out_name: "muxed.avi", | |
332 | mux_reg, enc_reg: RegisteredEncoders::new(), | |
333 | }; | |
334 | test_remuxing(&dec_config, &enc_config);*/ | |
335 | test_remuxing_md5(&dec_config, "avi", &mux_reg, | |
336 | [0xa0fb0e47, 0x412e24dd, 0x6b89711c, 0x276fb799]); | |
5cd2f537 KS |
337 | } |
338 | } |