1use crate::boxes::{Av1CBox, ClliBox, ColrBox, MdcvBox};
7
8#[derive(Clone)]
10#[non_exhaustive]
11pub struct AnimFrame<'a> {
12 pub color: &'a [u8],
14 pub alpha: Option<&'a [u8]>,
16 pub duration: u32,
18 pub is_sync: bool,
20}
21
22impl<'a> AnimFrame<'a> {
23 pub fn new(color: &'a [u8], duration: u32) -> Self {
25 Self { color, alpha: None, duration, is_sync: false }
26 }
27
28 pub fn with_alpha(mut self, alpha: &'a [u8]) -> Self {
30 self.alpha = Some(alpha);
31 self
32 }
33
34 pub fn with_sync(mut self, is_sync: bool) -> Self {
36 self.is_sync = is_sync;
37 self
38 }
39}
40
41pub struct AnimatedImage {
46 timescale: u32,
47 loop_count: u32,
48 color_config: Av1CBox,
49 alpha_config: Option<Av1CBox>,
50 colr: Option<ColrBox>,
51 clli: Option<ClliBox>,
52 mdcv: Option<MdcvBox>,
53}
54
55impl Default for AnimatedImage {
56 fn default() -> Self { Self::new() }
57}
58
59impl AnimatedImage {
60 pub fn new() -> Self {
62 Self {
63 timescale: 1000,
64 loop_count: 0,
65 color_config: Av1CBox::default(),
66 alpha_config: None,
67 colr: None,
68 clli: None,
69 mdcv: None,
70 }
71 }
72
73 pub fn set_timescale(&mut self, timescale: u32) -> &mut Self { self.timescale = timescale; self }
75 pub fn set_loop_count(&mut self, loop_count: u32) -> &mut Self { self.loop_count = loop_count; self }
77 pub fn set_color_config(&mut self, config: Av1CBox) -> &mut Self { self.color_config = config; self }
79 pub fn set_alpha_config(&mut self, config: Av1CBox) -> &mut Self { self.alpha_config = Some(config); self }
81 pub fn set_colr(&mut self, colr: ColrBox) -> &mut Self { self.colr = Some(colr); self }
83 pub fn set_clli(&mut self, clli: ClliBox) -> &mut Self { self.clli = Some(clli); self }
85 pub fn set_mdcv(&mut self, mdcv: MdcvBox) -> &mut Self { self.mdcv = Some(mdcv); self }
87
88 pub fn serialize(&self, width: u32, height: u32, frames: &[AnimFrame<'_>],
90 color_seq_header: &[u8], alpha_seq_header: Option<&[u8]>) -> Vec<u8> {
91 let has_alpha = frames.iter().any(|f| f.alpha.is_some())
92 && self.alpha_config.is_some()
93 && alpha_seq_header.is_some();
94
95 let total_duration: u64 = frames.iter().map(|f| u64::from(f.duration)).sum();
96 let durations: Vec<u32> = frames.iter().map(|f| f.duration).collect();
97 let color_frames: Vec<&[u8]> = frames.iter().map(|f| f.color).collect();
98 let alpha_frames: Vec<&[u8]> = if has_alpha {
99 frames.iter().map(|f| f.alpha.unwrap_or(&[])).collect()
100 } else {
101 Vec::new()
102 };
103 let sync_indices: Vec<u32> = frames.iter().enumerate()
104 .filter(|(_, f)| f.is_sync)
105 .map(|(i, _)| (i + 1) as u32) .collect();
107
108 let next_track_id = if has_alpha { 3 } else { 2 };
109
110 let mut out = Vec::new();
111
112 write_ftyp(&mut out);
114
115 write_meta(
117 &mut out,
118 width,
119 height,
120 color_seq_header,
121 color_frames.first().map(|f| f.len() as u32).unwrap_or(0),
122 &self.color_config,
123 self.colr.as_ref(),
124 self.clli.as_ref(),
125 self.mdcv.as_ref(),
126 );
127
128 let moov_pos = begin_box(&mut out, b"moov");
130 write_mvhd(&mut out, self.timescale, total_duration, next_track_id);
131 write_track(
132 &mut out, 1, width, height,
133 self.timescale, total_duration,
134 &color_frames, &durations, &sync_indices,
135 color_seq_header, &self.color_config,
136 false,
137 );
138 if has_alpha {
139 let alpha_seq = alpha_seq_header.unwrap();
140 let alpha_cfg = self.alpha_config.as_ref().unwrap();
141 write_track(
142 &mut out, 2, width, height,
143 self.timescale, total_duration,
144 &alpha_frames, &durations, &sync_indices,
145 alpha_seq, alpha_cfg,
146 true,
147 );
148 }
149 end_box(&mut out, moov_pos);
150
151 let mdat_pos = begin_box(&mut out, b"mdat");
153 let mdat_data_start = out.len();
154 for frame in &color_frames {
155 out.extend_from_slice(frame);
156 }
157 let alpha_data_start = out.len();
158 for frame in &alpha_frames {
159 out.extend_from_slice(frame);
160 }
161 end_box(&mut out, mdat_pos);
162
163 if has_alpha {
165 patch_offset_placeholders(&mut out, &[mdat_data_start as u32, alpha_data_start as u32], mdat_data_start as u32);
166 } else {
167 patch_offset_placeholders(&mut out, &[mdat_data_start as u32], mdat_data_start as u32);
168 }
169
170 out
171 }
172}
173
174fn write_u16(out: &mut Vec<u8>, v: u16) {
177 out.extend_from_slice(&v.to_be_bytes());
178}
179
180fn write_u32(out: &mut Vec<u8>, v: u32) {
181 out.extend_from_slice(&v.to_be_bytes());
182}
183
184fn write_u64(out: &mut Vec<u8>, v: u64) {
185 out.extend_from_slice(&v.to_be_bytes());
186}
187
188fn begin_box(out: &mut Vec<u8>, box_type: &[u8; 4]) -> usize {
190 let pos = out.len();
191 write_u32(out, 0); out.extend_from_slice(box_type);
193 pos
194}
195
196fn end_box(out: &mut [u8], pos: usize) {
198 let size = (out.len() - pos) as u32;
199 out[pos..pos + 4].copy_from_slice(&size.to_be_bytes());
200}
201
202fn write_fullbox(out: &mut Vec<u8>, version: u8, flags: u32) {
203 out.push(version);
204 out.push((flags >> 16) as u8);
205 out.push((flags >> 8) as u8);
206 out.push(flags as u8);
207}
208
209const STCO_PLACEHOLDER: u32 = 0xDEAD_BEEF;
210const ILOC_PLACEHOLDER: u32 = 0xDEAD_BEE0;
211
212fn write_ftyp(out: &mut Vec<u8>) {
215 let pos = begin_box(out, b"ftyp");
216 out.extend_from_slice(b"avis"); write_u32(out, 0); out.extend_from_slice(b"avis"); out.extend_from_slice(b"avif");
220 out.extend_from_slice(b"mif1");
221 out.extend_from_slice(b"miaf");
222 out.extend_from_slice(b"iso8");
223 end_box(out, pos);
224}
225
226#[allow(clippy::too_many_arguments)]
227fn write_meta(
228 out: &mut Vec<u8>,
229 width: u32,
230 height: u32,
231 seq_header: &[u8],
232 first_frame_len: u32,
233 av1c: &Av1CBox,
234 colr: Option<&ColrBox>,
235 clli: Option<&ClliBox>,
236 mdcv: Option<&MdcvBox>,
237) {
238 let meta_pos = begin_box(out, b"meta");
239 write_fullbox(out, 0, 0);
240
241 {
243 let pos = begin_box(out, b"hdlr");
244 write_fullbox(out, 0, 0);
245 write_u32(out, 0); out.extend_from_slice(b"pict");
247 out.extend_from_slice(&[0u8; 12]); out.push(0); end_box(out, pos);
250 }
251
252 {
254 let pos = begin_box(out, b"pitm");
255 write_fullbox(out, 0, 0);
256 write_u16(out, 1); end_box(out, pos);
258 }
259
260 {
262 let pos = begin_box(out, b"iloc");
263 write_fullbox(out, 0, 0);
264 out.push(0x44); out.push(0x00); write_u16(out, 1); write_u16(out, 1); write_u16(out, 0); write_u16(out, 1); write_u32(out, ILOC_PLACEHOLDER); write_u32(out, first_frame_len); end_box(out, pos);
273 }
274
275 {
277 let iinf_pos = begin_box(out, b"iinf");
278 write_fullbox(out, 0, 0);
279 write_u16(out, 1); let infe_pos = begin_box(out, b"infe");
282 write_fullbox(out, 2, 0);
283 write_u16(out, 1); write_u16(out, 0); out.extend_from_slice(b"av01");
286 out.push(0); end_box(out, infe_pos);
288
289 end_box(out, iinf_pos);
290 }
291
292 {
294 let iprp_pos = begin_box(out, b"iprp");
295
296 {
298 let ipco_pos = begin_box(out, b"ipco");
299
300 {
302 let pos = begin_box(out, b"ispe");
303 write_fullbox(out, 0, 0);
304 write_u32(out, width);
305 write_u32(out, height);
306 end_box(out, pos);
307 }
308
309 write_av1c_box(out, av1c, seq_header);
311
312 {
314 let pos = begin_box(out, b"pixi");
315 write_fullbox(out, 0, 0);
316 if av1c.monochrome {
317 out.push(1); out.push(bit_depth_from_av1c(av1c));
319 } else {
320 out.push(3); let depth = bit_depth_from_av1c(av1c);
322 out.push(depth);
323 out.push(depth);
324 out.push(depth);
325 }
326 end_box(out, pos);
327 }
328
329 if let Some(colr) = colr
331 && *colr != ColrBox::default() {
332 write_colr_nclx(out, colr);
333 }
334
335 if let Some(clli) = clli {
337 write_clli(out, clli);
338 }
339
340 if let Some(mdcv) = mdcv {
342 write_mdcv(out, mdcv);
343 }
344
345 end_box(out, ipco_pos);
346 }
347
348 {
350 let pos = begin_box(out, b"ipma");
351 write_fullbox(out, 0, 0);
352 write_u32(out, 1); write_u16(out, 1); let mut assoc_count: u8 = 3;
356 let has_colr = colr.is_some_and(|c| *c != ColrBox::default());
357 if has_colr { assoc_count += 1; }
358 if clli.is_some() { assoc_count += 1; }
359 if mdcv.is_some() { assoc_count += 1; }
360 out.push(assoc_count);
361 out.push(0x01); out.push(0x82); out.push(0x03); let mut next_prop = 4u8;
365 if has_colr {
366 out.push(next_prop);
367 next_prop += 1;
368 }
369 if clli.is_some() {
370 out.push(next_prop);
371 next_prop += 1;
372 }
373 if mdcv.is_some() {
374 out.push(next_prop);
375 let _ = next_prop;
376 }
377 end_box(out, pos);
378 }
379
380 end_box(out, iprp_pos);
381 }
382
383 end_box(out, meta_pos);
384}
385
386fn write_mvhd(out: &mut Vec<u8>, timescale: u32, duration: u64, next_track_id: u32) {
387 let pos = begin_box(out, b"mvhd");
388 write_fullbox(out, 1, 0);
389 write_u64(out, 0); write_u64(out, 0); write_u32(out, timescale);
392 write_u64(out, duration);
393 write_u32(out, 0x0001_0000); write_u16(out, 0x0100); out.extend_from_slice(&[0u8; 10]); for &v in &[0x0001_0000u32, 0, 0, 0, 0x0001_0000, 0, 0, 0, 0x4000_0000] {
398 write_u32(out, v);
399 }
400 out.extend_from_slice(&[0u8; 24]); write_u32(out, next_track_id);
402 end_box(out, pos);
403}
404
405#[allow(clippy::too_many_arguments)]
406fn write_track(
407 out: &mut Vec<u8>,
408 track_id: u32,
409 width: u32,
410 height: u32,
411 timescale: u32,
412 duration: u64,
413 frames: &[&[u8]],
414 durations: &[u32],
415 sync_indices: &[u32],
416 seq_header: &[u8],
417 av1c: &Av1CBox,
418 is_alpha: bool,
419) {
420 let trak_pos = begin_box(out, b"trak");
421
422 {
424 let pos = begin_box(out, b"tkhd");
425 let flags = if is_alpha { 1 } else { 3 }; write_fullbox(out, 1, flags);
427 write_u64(out, 0); write_u64(out, 0); write_u32(out, track_id);
430 write_u32(out, 0); write_u64(out, duration);
432 out.extend_from_slice(&[0u8; 8]); write_u16(out, 0); write_u16(out, 0); write_u16(out, 0); write_u16(out, 0); for &v in &[0x0001_0000u32, 0, 0, 0, 0x0001_0000, 0, 0, 0, 0x4000_0000] {
438 write_u32(out, v);
439 }
440 write_u32(out, width << 16);
441 write_u32(out, height << 16);
442 end_box(out, pos);
443 }
444
445 {
447 let mdia_pos = begin_box(out, b"mdia");
448
449 {
451 let pos = begin_box(out, b"mdhd");
452 write_fullbox(out, 1, 0);
453 write_u64(out, 0); write_u64(out, 0); write_u32(out, timescale);
456 write_u64(out, duration);
457 write_u16(out, 0x55C4); write_u16(out, 0);
459 end_box(out, pos);
460 }
461
462 {
464 let pos = begin_box(out, b"hdlr");
465 write_fullbox(out, 0, 0);
466 write_u32(out, 0);
467 if is_alpha {
468 out.extend_from_slice(b"auxv");
469 } else {
470 out.extend_from_slice(b"pict");
471 }
472 out.extend_from_slice(&[0u8; 12]);
473 out.extend_from_slice(if is_alpha { b"Alpha\0" } else { b"Color\0" });
474 end_box(out, pos);
475 }
476
477 {
479 let minf_pos = begin_box(out, b"minf");
480
481 {
483 let pos = begin_box(out, b"vmhd");
484 write_fullbox(out, 0, 1);
485 out.extend_from_slice(&[0u8; 8]); end_box(out, pos);
487 }
488
489 {
491 let dinf_pos = begin_box(out, b"dinf");
492 let dref_pos = begin_box(out, b"dref");
493 write_fullbox(out, 0, 0);
494 write_u32(out, 1);
495 let url_pos = begin_box(out, b"url ");
496 write_fullbox(out, 0, 1); end_box(out, url_pos);
498 end_box(out, dref_pos);
499 end_box(out, dinf_pos);
500 }
501
502 {
504 let stbl_pos = begin_box(out, b"stbl");
505
506 {
508 let pos = begin_box(out, b"stsd");
509 write_fullbox(out, 0, 0);
510 write_u32(out, 1); let av01_pos = begin_box(out, b"av01");
513 out.extend_from_slice(&[0u8; 6]); write_u16(out, 1); write_u16(out, 0); write_u16(out, 0); out.extend_from_slice(&[0u8; 12]); write_u16(out, width as u16);
519 write_u16(out, height as u16);
520 write_u32(out, 0x0048_0000); write_u32(out, 0x0048_0000); write_u32(out, 0); write_u16(out, 1); out.extend_from_slice(&[0u8; 32]); write_u16(out, 0x0018); out.extend_from_slice(&0xFFFFu16.to_be_bytes()); write_av1c_box(out, av1c, seq_header);
530
531 end_box(out, av01_pos);
532 end_box(out, pos);
533 }
534
535 {
537 let pos = begin_box(out, b"stts");
538 write_fullbox(out, 0, 0);
539 let mut entries: Vec<(u32, u32)> = Vec::new();
540 for &d in durations {
541 if let Some(last) = entries.last_mut()
542 && last.1 == d {
543 last.0 += 1;
544 continue;
545 }
546 entries.push((1, d));
547 }
548 write_u32(out, entries.len() as u32);
549 for (count, delta) in &entries {
550 write_u32(out, *count);
551 write_u32(out, *delta);
552 }
553 end_box(out, pos);
554 }
555
556 {
558 let pos = begin_box(out, b"stsc");
559 write_fullbox(out, 0, 0);
560 write_u32(out, 1);
561 write_u32(out, 1); write_u32(out, frames.len() as u32); write_u32(out, 1); end_box(out, pos);
565 }
566
567 {
569 let pos = begin_box(out, b"stsz");
570 write_fullbox(out, 0, 0);
571 write_u32(out, 0); write_u32(out, frames.len() as u32);
573 for frame in frames {
574 write_u32(out, frame.len() as u32);
575 }
576 end_box(out, pos);
577 }
578
579 {
581 let pos = begin_box(out, b"stco");
582 write_fullbox(out, 0, 0);
583 write_u32(out, 1); write_u32(out, STCO_PLACEHOLDER);
585 end_box(out, pos);
586 }
587
588 {
590 let pos = begin_box(out, b"stss");
591 write_fullbox(out, 0, 0);
592 write_u32(out, sync_indices.len() as u32);
593 for &idx in sync_indices {
594 write_u32(out, idx);
595 }
596 end_box(out, pos);
597 }
598
599 end_box(out, stbl_pos);
600 }
601
602 end_box(out, minf_pos);
603 }
604
605 end_box(out, mdia_pos);
606 }
607
608 if is_alpha {
610 let tref_pos = begin_box(out, b"tref");
611 let auxl_pos = begin_box(out, b"auxl");
612 write_u32(out, 1); end_box(out, auxl_pos);
614 end_box(out, tref_pos);
615 }
616
617 end_box(out, trak_pos);
618}
619
620fn write_av1c_box(out: &mut Vec<u8>, av1c: &Av1CBox, seq_header: &[u8]) {
623 let pos = begin_box(out, b"av1C");
624 out.push(0x81); let byte1 = (av1c.seq_profile << 5) | av1c.seq_level_idx_0;
627 let byte2 =
628 u8::from(av1c.seq_tier_0) << 7
629 | u8::from(av1c.high_bitdepth) << 6
630 | u8::from(av1c.twelve_bit) << 5
631 | u8::from(av1c.monochrome) << 4
632 | u8::from(av1c.chroma_subsampling_x) << 3
633 | u8::from(av1c.chroma_subsampling_y) << 2
634 | av1c.chroma_sample_position;
635
636 out.push(byte1);
637 out.push(byte2);
638 out.push(0x00); out.extend_from_slice(seq_header);
640 end_box(out, pos);
641}
642
643fn bit_depth_from_av1c(av1c: &Av1CBox) -> u8 {
644 if av1c.twelve_bit { 12 } else if av1c.high_bitdepth { 10 } else { 8 }
645}
646
647fn write_colr_nclx(out: &mut Vec<u8>, colr: &ColrBox) {
648 let pos = begin_box(out, b"colr");
649 out.extend_from_slice(b"nclx");
650 write_u16(out, colr.color_primaries as u16);
651 write_u16(out, colr.transfer_characteristics as u16);
652 write_u16(out, colr.matrix_coefficients as u16);
653 out.push(if colr.full_range_flag { 1 << 7 } else { 0 });
654 end_box(out, pos);
655}
656
657fn write_clli(out: &mut Vec<u8>, clli: &ClliBox) {
658 let pos = begin_box(out, b"clli");
659 write_u16(out, clli.max_content_light_level);
660 write_u16(out, clli.max_pic_average_light_level);
661 end_box(out, pos);
662}
663
664fn write_mdcv(out: &mut Vec<u8>, mdcv: &MdcvBox) {
665 let pos = begin_box(out, b"mdcv");
666 for &(x, y) in &mdcv.primaries {
667 write_u16(out, x);
668 write_u16(out, y);
669 }
670 write_u16(out, mdcv.white_point.0);
671 write_u16(out, mdcv.white_point.1);
672 write_u32(out, mdcv.max_luminance);
673 write_u32(out, mdcv.min_luminance);
674 end_box(out, pos);
675}
676
677fn patch_offset_placeholders(out: &mut [u8], stco_offsets: &[u32], iloc_offset: u32) {
679 let stco_placeholder = STCO_PLACEHOLDER.to_be_bytes();
680 let iloc_placeholder = ILOC_PLACEHOLDER.to_be_bytes();
681 let mut stco_idx = 0;
682 let mut i = 0;
683 while i + 4 <= out.len() {
684 if stco_idx < stco_offsets.len() && out[i..i + 4] == stco_placeholder {
685 out[i..i + 4].copy_from_slice(&stco_offsets[stco_idx].to_be_bytes());
686 stco_idx += 1;
687 i += 4;
688 } else if out[i..i + 4] == iloc_placeholder {
689 out[i..i + 4].copy_from_slice(&iloc_offset.to_be_bytes());
690 i += 4;
691 } else {
692 i += 1;
693 }
694 }
695}
696
697#[cfg(test)]
698mod tests {
699 use super::*;
700
701 fn basic_av1c() -> Av1CBox {
702 Av1CBox {
703 seq_profile: 0,
704 seq_level_idx_0: 4,
705 seq_tier_0: false,
706 high_bitdepth: false,
707 twelve_bit: false,
708 monochrome: false,
709 chroma_subsampling_x: true,
710 chroma_subsampling_y: true,
711 chroma_sample_position: 0,
712 }
713 }
714
715 fn mono_av1c() -> Av1CBox {
716 Av1CBox {
717 seq_profile: 0,
718 seq_level_idx_0: 4,
719 seq_tier_0: false,
720 high_bitdepth: false,
721 twelve_bit: false,
722 monochrome: true,
723 chroma_subsampling_x: true,
724 chroma_subsampling_y: true,
725 chroma_sample_position: 0,
726 }
727 }
728
729 #[test]
730 fn serialize_color_only() {
731 let frames = [
732 AnimFrame::new(b"frame1color", 100).with_sync(true),
733 AnimFrame::new(b"frame2color", 200),
734 ];
735 let mut image = AnimatedImage::new();
736 image.set_color_config(basic_av1c());
737 let avif = image.serialize(64, 64, &frames, b"seqhdr", None);
738
739 assert_eq!(&avif[4..8], b"ftyp");
741 assert_eq!(&avif[8..12], b"avis");
742
743 let mdat_str = b"mdat";
745 assert!(avif.windows(4).any(|w| w == mdat_str));
746
747 assert!(avif.windows(b"frame1color".len()).any(|w| w == b"frame1color"));
749 assert!(avif.windows(b"frame2color".len()).any(|w| w == b"frame2color"));
750
751 let parser = zenavif_parse::AvifParser::from_bytes(&avif).unwrap();
753 let info = parser.animation_info().expect("should have animation info");
754 assert_eq!(info.timescale, 1000);
755 assert_eq!(info.frame_count, 2);
756 }
757
758 #[test]
759 fn serialize_with_alpha() {
760 let frames = [
761 AnimFrame::new(b"c1", 500).with_alpha(b"a1").with_sync(true),
762 AnimFrame::new(b"c2", 500).with_alpha(b"a2"),
763 ];
764 let mut image = AnimatedImage::new();
765 image.set_color_config(basic_av1c());
766 image.set_alpha_config(mono_av1c());
767 let avif = image.serialize(32, 32, &frames, b"colseq", Some(b"alphaseq"));
768
769 assert_eq!(&avif[4..8], b"ftyp");
770 assert!(avif.windows(2).any(|w| w == b"c1"));
771 assert!(avif.windows(2).any(|w| w == b"a1"));
772 assert!(avif.windows(2).any(|w| w == b"c2"));
773 assert!(avif.windows(2).any(|w| w == b"a2"));
774
775 let parser = zenavif_parse::AvifParser::from_bytes(&avif).unwrap();
776 let info = parser.animation_info().expect("should have animation info");
777 assert_eq!(info.frame_count, 2);
778 }
779
780 #[test]
781 fn frame_durations_roundtrip() {
782 let frames = [
783 AnimFrame::new(b"f1", 100).with_sync(true),
784 AnimFrame::new(b"f2", 200),
785 AnimFrame::new(b"f3", 300),
786 ];
787 let mut image = AnimatedImage::new();
788 image.set_color_config(basic_av1c());
789 let avif = image.serialize(16, 16, &frames, b"seq", None);
790 let parser = zenavif_parse::AvifParser::from_bytes(&avif).unwrap();
791 let info = parser.animation_info().expect("animation info");
792 assert_eq!(info.frame_count, 3);
793 assert_eq!(info.timescale, 1000);
794 }
795}