avimux: set palette change flag in stream header if applicable
[nihav.git] / nihav-commonfmt / src / muxers / avi.rs
CommitLineData
5cd2f537
KS
1use nihav_core::muxers::*;
2use nihav_registry::register::*;
3
4#[derive(Clone,Copy)]
5struct IdxEntry {
6 stream: u32,
7 stype: StreamType,
8 key: bool,
9 pos: u32,
10 len: u32,
11}
12
13#[derive(Clone,Copy)]
14struct AVIStream {
15 strh_pos: u64,
16 nframes: u32,
17 is_video: bool,
c8db9313 18 max_size: u32,
04fafc56 19 pal_change: bool,
5cd2f537
KS
20}
21
22struct AVIMuxer<'a> {
23 bw: &'a mut ByteWriter<'a>,
24 index: Vec<IdxEntry>,
25 video_str: Option<usize>,
26 video_id: u32,
27 data_pos: u64,
28 stream_info: Vec<AVIStream>,
a1613eee 29 pal_pos: Vec<u32>,
5cd2f537
KS
30}
31
32impl<'a> AVIMuxer<'a> {
33 fn new(bw: &'a mut ByteWriter<'a>) -> Self {
34 Self {
35 bw,
36 index: Vec::new(),
37 video_str: None,
38 video_id: 0,
39 data_pos: 0,
40 stream_info: Vec::with_capacity(2),
a1613eee 41 pal_pos: Vec::with_capacity(2),
5cd2f537
KS
42 }
43 }
44}
45
46fn patch_size(bw: &mut ByteWriter, pos: u64) -> MuxerResult<()> {
47 let size = bw.tell() - pos;
48 bw.seek(SeekFrom::Current(-((size + 4) as i64)))?;
49 bw.write_u32le(size as u32)?;
50 bw.seek(SeekFrom::End(0))?;
51 Ok(())
52}
53
54fn write_chunk_hdr(bw: &mut ByteWriter, stype: StreamType, str_no: u32) -> MuxerResult<()> {
55 bw.write_byte(b'0' + ((str_no / 10) as u8))?;
56 bw.write_byte(b'0' + ((str_no % 10) as u8))?;
57 match stype {
58 StreamType::Video => { bw.write_buf(b"dc")?; },
59 StreamType::Audio => { bw.write_buf(b"wb")?; },
60 StreamType::Subtitles => { bw.write_buf(b"tx")?; },
61 _ => return Err(MuxerError::UnsupportedFormat),
62 };
63 Ok(())
64}
65
66impl<'a> MuxCore<'a> for AVIMuxer<'a> {
61cab15b
KS
67 #[allow(clippy::unreadable_literal)]
68 #[allow(clippy::cast_lossless)]
5cd2f537
KS
69 fn create(&mut self, strmgr: &StreamManager) -> MuxerResult<()> {
70 if strmgr.get_num_streams() == 0 {
71 return Err(MuxerError::InvalidArgument);
72 }
73 if strmgr.get_num_streams() > 99 {
74 return Err(MuxerError::UnsupportedFormat);
75 }
405cec9e
KS
76 for (str_no, strm) in strmgr.iter().enumerate() {
77 if strm.get_media_type() == StreamType::Video {
5cd2f537 78 self.video_str = Some(str_no);
405cec9e 79 self.video_id = strm.id;
5cd2f537
KS
80 break;
81 }
82 }
83 let (vinfo, tb_num, tb_den) = if let Some(str_id) = self.video_str {
84 let vstr = strmgr.get_stream(str_id).unwrap();
85 (vstr.get_info(), vstr.tb_num, vstr.tb_den)
86 } else {
87 (NACodecInfo::new_dummy(), 0, 1)
88 };
89 let hdrl_pos = self.bw.tell() + 20;
90 self.bw.write_buf(b"RIFF\0\0\0\0AVI LIST\0\0\0\0hdrlavih")?;
91 self.bw.write_u32le(56)?; // avih size
92 let ms_per_frame = NATimeInfo::ts_to_time(1, 1000000, tb_num, tb_den);
93 self.bw.write_u32le(ms_per_frame as u32)?;
94 self.bw.write_u32le(0)?; // max transfer rate
95 self.bw.write_u32le(0)?; // padding granularity
96 self.bw.write_u32le(0)?; // flags
97 self.bw.write_u32le(0)?; // total frames
98 self.bw.write_u32le(0)?; // initial frames
99 self.bw.write_u32le(strmgr.get_num_streams() as u32)?;
100 self.bw.write_u32le(0)?; // suggested buffer size
101 if let NACodecTypeInfo::Video(ref vinfo) = vinfo.get_properties() {
102 self.bw.write_u32le(vinfo.width as u32)?;
103 self.bw.write_u32le(vinfo.height as u32)?;
104 } else {
105 self.bw.write_u32le(0)?;
106 self.bw.write_u32le(0)?;
107 }
108 self.bw.write_u32le(0)?; // reserved
109 self.bw.write_u32le(0)?; // reserved
110 self.bw.write_u32le(0)?; // reserved
111 self.bw.write_u32le(0)?; // reserved
112
37952415 113 self.pal_pos.clear();
a1613eee 114 self.pal_pos.resize(strmgr.get_num_streams(), 0);
405cec9e 115 for (strno, strm) in strmgr.iter().enumerate() {
5cd2f537
KS
116 let strl_pos = self.bw.tell() + 8;
117 self.bw.write_buf(b"LIST\0\0\0\0strlstrh")?;
118 self.bw.write_u32le(56)?; // strh size
119
405cec9e 120 match strm.get_media_type() {
5cd2f537
KS
121 StreamType::Video => {
122 self.bw.write_buf(b"vids")?;
405cec9e 123 let fcc = find_avi_fourcc(strm.get_info().get_name());
5cd2f537
KS
124 if fcc.is_none() {
125 return Err(MuxerError::UnsupportedFormat);
126 }
127 self.bw.write_buf(&fcc.unwrap_or([0; 4]))?;
405cec9e 128 let vinfo = strm.get_info().get_properties().get_video_info().unwrap();
5cd2f537
KS
129 if vinfo.width >= (1 << 16) || vinfo.height >= (1 << 16) {
130 return Err(MuxerError::UnsupportedFormat);
131 }
132 },
133 StreamType::Audio => {
134 self.bw.write_buf(b"auds")?;
135 self.bw.write_u32le(0)?;
136 },
137 StreamType::Subtitles => {
138 self.bw.write_buf(b"txts")?;
139 self.bw.write_u32le(0)?;
140 },
141 _ => return Err(MuxerError::UnsupportedFormat),
142 };
143 self.stream_info.push(AVIStream {
144 strh_pos: self.bw.tell(),
405cec9e 145 is_video: strm.get_media_type() == StreamType::Video,
5cd2f537
KS
146 nframes: 0,
147 max_size: 0,
04fafc56 148 pal_change: false,
5cd2f537
KS
149 });
150
151 self.bw.write_u32le(0)?; // flags
152 self.bw.write_u16le(0)?; // priority
153 self.bw.write_u16le(0)?; // language
154 self.bw.write_u32le(0)?; // initial frames
405cec9e
KS
155 self.bw.write_u32le(strm.tb_num)?;
156 self.bw.write_u32le(strm.tb_den)?;
5cd2f537
KS
157 self.bw.write_u32le(0)?; // start
158 self.bw.write_u32le(0)?; // length
159 self.bw.write_u32le(0)?; // suggested buffer size
160 self.bw.write_u32le(0)?; // quality
161 self.bw.write_u32le(0)?; // sample_size
162 self.bw.write_u16le(0)?; // x
163 self.bw.write_u16le(0)?; // y
164 self.bw.write_u16le(0)?; // w
165 self.bw.write_u16le(0)?; // h
166
167 self.bw.write_buf(b"strf")?;
168 self.bw.write_u32le(0)?;
169 let strf_pos = self.bw.tell();
405cec9e 170 match strm.get_media_type() {
5cd2f537 171 StreamType::Video => {
405cec9e 172 let vinfo = strm.get_info().get_properties().get_video_info().unwrap();
5cd2f537
KS
173 let hdr_pos = self.bw.tell();
174 self.bw.write_u32le(0)?;
175 self.bw.write_u32le(vinfo.width as u32)?;
e5b5248d 176 self.bw.write_u32le(vinfo.height as u32)?;
c905da8d
KS
177 if !vinfo.format.palette {
178 self.bw.write_u16le(vinfo.format.components as u16)?;
179 self.bw.write_u16le(vinfo.format.get_total_depth() as u16)?;
180 } else {
181 self.bw.write_u16le(1)?;
182 self.bw.write_u16le(8)?;
183 }
405cec9e 184 let fcc = find_avi_fourcc(strm.get_info().get_name());
5cd2f537
KS
185 if fcc.is_none() {
186 return Err(MuxerError::UnsupportedFormat);
187 }
188 self.bw.write_buf(&fcc.unwrap_or([0; 4]))?;
189 self.bw.write_u32le(0)?; // image size
190 self.bw.write_u32le(0)?; // x dpi
191 self.bw.write_u32le(0)?; // y dpi
192 if vinfo.format.palette {
a1613eee 193 self.bw.write_u32le(256)?; // total colors
5cd2f537 194 self.bw.write_u32le(0)?; // important colors
a1613eee
KS
195 self.pal_pos[strno] = self.bw.tell() as u32;
196 for _ in 0..256 {
197 self.bw.write_u32le(0)?;
198 }
5cd2f537
KS
199 } else {
200 self.bw.write_u32le(0)?; // total colors
201 self.bw.write_u32le(0)?; // important colors
202 }
405cec9e 203 if let Some(ref edata) = strm.get_info().get_extradata() {
5cd2f537
KS
204 self.bw.write_buf(edata.as_slice())?;
205 }
206 let bisize = self.bw.tell() - hdr_pos;
207 self.bw.seek(SeekFrom::Current(-(bisize as i64)))?;
208 self.bw.write_u32le(bisize as u32)?;
209 self.bw.seek(SeekFrom::End(0))?;
210 },
211 StreamType::Audio => {
405cec9e
KS
212 let ainfo = strm.get_info().get_properties().get_audio_info().unwrap();
213 let twocc = find_wav_twocc(strm.get_info().get_name());
5cd2f537
KS
214 if twocc.is_none() {
215 return Err(MuxerError::UnsupportedFormat);
216 }
586ea187 217 self.bw.write_u16le(twocc.unwrap_or(0))?;
5cd2f537
KS
218 self.bw.write_u16le(ainfo.channels as u16)?;
219 self.bw.write_u32le(ainfo.sample_rate)?;
220 self.bw.write_u32le(0)?; // avg bytes per second
221 self.bw.write_u16le(ainfo.block_len as u16)?;
222 self.bw.write_u16le(ainfo.format.bits as u16)?;
405cec9e 223 if let Some(ref edata) = strm.get_info().get_extradata() {
5cd2f537
KS
224 self.bw.write_buf(edata.as_slice())?;
225 }
226 },
227 StreamType::Subtitles => {
405cec9e 228 if let Some(ref edata) = strm.get_info().get_extradata() {
5cd2f537
KS
229 self.bw.write_buf(edata.as_slice())?;
230 }
231 },
232 _ => unreachable!(),
233 };
6f263099
KS
234 patch_size(self.bw, strf_pos)?;
235 patch_size(self.bw, strl_pos)?;
5cd2f537 236 }
6f263099 237 patch_size(self.bw, hdrl_pos)?;
5cd2f537
KS
238
239 self.data_pos = self.bw.tell() + 8;
240 self.bw.write_buf(b"LIST\0\0\0\0movi")?;
241
242 Ok(())
243 }
244 fn mux_frame(&mut self, _strmgr: &StreamManager, pkt: NAPacket) -> MuxerResult<()> {
245 if self.data_pos == 0 {
246 return Err(MuxerError::NotCreated);
247 }
817e4872
KS
248 let stream = pkt.get_stream();
249 let str_num = stream.get_num();
5cd2f537
KS
250 if str_num > 99 || str_num >= self.stream_info.len() {
251 return Err(MuxerError::UnsupportedFormat);
252 }
253
254 let chunk_len = pkt.get_buffer().len() as u32;
255
a1613eee
KS
256 if self.pal_pos[str_num] != 0 {
257 for sdata in pkt.side_data.iter() {
258 if let NASideData::Palette(_, ref pal) = sdata {
259 let cur_pos = self.bw.tell();
260 self.bw.seek(SeekFrom::Start(u64::from(self.pal_pos[str_num])))?;
bfe6df94
KS
261 for quad in pal.chunks(4) {
262 self.bw.write_byte(quad[2])?;
263 self.bw.write_byte(quad[1])?;
264 self.bw.write_byte(quad[0])?;
265 self.bw.write_byte(0)?;
266 }
a1613eee
KS
267 self.bw.seek(SeekFrom::Start(cur_pos))?;
268 self.pal_pos[str_num] = 0;
269 break;
270 }
271 }
272 } else {
273 for sdata in pkt.side_data.iter() {
274 if let NASideData::Palette(true, ref pal) = sdata {
275 //todo search for changed region
276 let start_clr = 0usize;
277 let end_clr = 256usize;
278 if start_clr < end_clr {
279 let chunk_len = ((end_clr - start_clr) as u32) * 4 + 4;
280 self.bw.write_byte(b'0' + ((str_num / 10) as u8))?;
281 self.bw.write_byte(b'0' + ((str_num % 10) as u8))?;
282 self.bw.write_buf(b"pc")?;
283 self.bw.write_u32le(chunk_len)?;
284 self.bw.write_byte(start_clr as u8)?;
285 self.bw.write_byte((end_clr - start_clr) as u8)?;
286 self.bw.write_u16le(0)?; //flags
287 self.bw.write_buf(&pal[start_clr * 4..end_clr * 4])?;
04fafc56 288 self.stream_info[str_num].pal_change = true;
a1613eee
KS
289 }
290 }
291 }
292 }
293
5cd2f537
KS
294 self.stream_info[str_num].nframes += 1;
295 self.stream_info[str_num].max_size = self.stream_info[str_num].max_size.max(chunk_len);
5cd2f537
KS
296 self.index.push(IdxEntry {
297 stream: str_num as u32,
817e4872 298 stype: stream.get_media_type(),
5cd2f537
KS
299 key: pkt.keyframe,
300 pos: self.bw.tell() as u32,
301 len: chunk_len });
817e4872 302 write_chunk_hdr(self.bw, stream.get_media_type(), str_num as u32)?;
5cd2f537
KS
303 self.bw.write_u32le(chunk_len)?;
304 self.bw.write_buf(pkt.get_buffer().as_slice())?;
77d52ff6
KS
305 if (self.bw.tell() & 1) != 0 {
306 self.bw.write_byte(0)?;
307 }
5cd2f537
KS
308 Ok(())
309 }
310 fn flush(&mut self) -> MuxerResult<()> {
311 Ok(())
312 }
313 fn end(&mut self) -> MuxerResult<()> {
6f263099 314 patch_size(self.bw, self.data_pos)?;
61cab15b 315 if !self.index.is_empty() {
5cd2f537
KS
316 self.bw.write_buf(b"idx1")?;
317 self.bw.write_u32le((self.index.len() * 16) as u32)?;
318 for item in self.index.iter() {
6f263099 319 write_chunk_hdr(self.bw, item.stype, item.stream)?;
5cd2f537
KS
320 self.bw.write_u32le(if item.key { 0x10 } else { 0 })?;
321 self.bw.write_u32le(item.pos)?;
322 self.bw.write_u32le(item.len)?;
323 }
324 }
6f263099 325 patch_size(self.bw, 8)?;
5cd2f537
KS
326 let mut max_frames = 0;
327 let mut max_size = 0;
328 for stri in self.stream_info.iter() {
329 max_frames = max_frames.max(stri.nframes);
330 max_size = max_size.max(stri.max_size);
04fafc56
KS
331 if stri.pal_change {
332 self.bw.seek(SeekFrom::Start(stri.strh_pos))?;
333 self.bw.write_u32le(0x00010000)?;
334 }
5cd2f537
KS
335 self.bw.seek(SeekFrom::Start(stri.strh_pos + 0x18))?;
336 self.bw.write_u32le(if stri.is_video { stri.nframes } else { 0 })?;
337 self.bw.write_u32le(stri.max_size)?;
338 }
339 self.bw.seek(SeekFrom::Start(0x30))?;
340 self.bw.write_u32le(max_frames)?;
341 self.bw.seek(SeekFrom::Current(8))?;
342 self.bw.write_u32le(max_size)?;
343 Ok(())
344 }
345}
346
dc80f48e
KS
347impl<'a> NAOptionHandler for AVIMuxer<'a> {
348 fn get_supported_options(&self) -> &[NAOptionDefinition] { &[] }
349 fn set_options(&mut self, _options: &[NAOption]) { }
350 fn query_option_value(&self, _name: &str) -> Option<NAValue> { None }
351}
352
5cd2f537
KS
353pub struct AVIMuxerCreator {}
354
355impl MuxerCreator for AVIMuxerCreator {
356 fn new_muxer<'a>(&self, bw: &'a mut ByteWriter<'a>) -> Box<dyn MuxCore<'a> + 'a> {
357 Box::new(AVIMuxer::new(bw))
358 }
359 fn get_name(&self) -> &'static str { "avi" }
f0081142 360 fn get_capabilities(&self) -> MuxerCapabilities { MuxerCapabilities::Universal }
5cd2f537
KS
361}
362
363#[cfg(test)]
364mod test {
68e5a4ca 365 use nihav_core::codecs::*;
5cd2f537 366 use nihav_core::demuxers::*;
68e5a4ca
KS
367 use nihav_core::muxers::*;
368 use nihav_codec_support::test::enc_video::*;
369 use crate::*;
5cd2f537
KS
370
371 #[test]
372 fn test_avi_muxer() {
373 let mut dmx_reg = RegisteredDemuxers::new();
374 generic_register_all_demuxers(&mut dmx_reg);
886cde48 375 //test sample: https://samples.mplayerhq.hu/V-codecs/RT21/320x240/laser05.avi
68e5a4ca
KS
376 let dec_config = DecoderTestParams {
377 demuxer: "avi",
378 in_name: "assets/Indeo/laser05.avi",
379 limit: None,
380 stream_type: StreamType::None,
381 dmx_reg, dec_reg: RegisteredDecoders::new(),
382 };
383 let mut mux_reg = RegisteredMuxers::new();
384 generic_register_all_muxers(&mut mux_reg);
385 /*let enc_config = EncoderTestParams {
386 muxer: "avi",
387 enc_name: "",
388 out_name: "muxed.avi",
389 mux_reg, enc_reg: RegisteredEncoders::new(),
390 };
391 test_remuxing(&dec_config, &enc_config);*/
392 test_remuxing_md5(&dec_config, "avi", &mux_reg,
393 [0xa0fb0e47, 0x412e24dd, 0x6b89711c, 0x276fb799]);
5cd2f537
KS
394 }
395}