1use crate::av1::encode::frame::{
21 CodedStill, DEFAULT_EFFORT, Picture, encode_still_with, qindex_for_quality,
22};
23use otf_pixels_core::{
24 ChromaSubsampling, EncodeOptions, Encoder, ImageDescriptor, PixelFormat, PixelsError, Result,
25 Sink,
26};
27
28const CICP: (u8, u8, u8) = (1, 13, 6);
30const MAX_DIMENSION: u32 = 65_536;
32
33#[derive(Debug)]
35pub struct AvifEncoder {
36 state: Option<State>,
37 options: EncodeOptions,
38 icc: Option<Vec<u8>>,
40 exif: Option<Vec<u8>>,
42 xmp: Option<Vec<u8>>,
44}
45
46#[derive(Debug)]
47struct State {
48 descriptor: ImageDescriptor,
49 pixels: Vec<u8>,
50 rows_written: u32,
51}
52
53impl Default for AvifEncoder {
54 fn default() -> Self {
55 Self::new()
56 }
57}
58
59impl AvifEncoder {
60 #[must_use]
62 pub fn new() -> Self {
63 Self::from_options(&EncodeOptions::default())
64 }
65
66 #[must_use]
68 pub const fn from_options(options: &EncodeOptions) -> Self {
69 Self {
70 state: None,
71 options: *options,
72 icc: None,
73 exif: None,
74 xmp: None,
75 }
76 }
77}
78
79impl Encoder for AvifEncoder {
80 fn set_icc_profile(&mut self, profile: Option<&[u8]>) -> Result<()> {
81 if self.state.is_some() {
82 return Err(PixelsError::invalid_argument(
83 "profile",
84 "the ICC profile must be set before write_header",
85 ));
86 }
87 self.icc = profile.map(<[u8]>::to_vec);
88 Ok(())
89 }
90
91 fn set_exif(&mut self, exif: Option<&[u8]>) -> Result<()> {
92 if self.state.is_some() {
93 return Err(PixelsError::invalid_argument(
94 "exif",
95 "the EXIF block must be set before write_header",
96 ));
97 }
98 self.exif = exif.map(<[u8]>::to_vec);
99 Ok(())
100 }
101
102 fn set_xmp(&mut self, xmp: Option<&[u8]>) -> Result<()> {
103 if self.state.is_some() {
104 return Err(PixelsError::invalid_argument(
105 "xmp",
106 "the XMP packet must be set before write_header",
107 ));
108 }
109 self.xmp = xmp.map(<[u8]>::to_vec);
110 Ok(())
111 }
112
113 fn write_header(&mut self, desc: &ImageDescriptor, _sink: &mut dyn Sink) -> Result<()> {
114 if self.state.is_some() {
115 return Err(PixelsError::invalid_argument(
116 "descriptor",
117 "write_header called more than once",
118 ));
119 }
120 if self.options.lossless {
121 return Err(PixelsError::unsupported(
122 "lossless AVIF encoding is not implemented; encode lossy or choose PNG or WebP",
123 ));
124 }
125 match desc.pixel {
126 PixelFormat::Gray8 | PixelFormat::GrayA8 | PixelFormat::Rgb8 | PixelFormat::Rgba8 => {}
127 other => {
128 return Err(PixelsError::unsupported(format!(
129 "AVIF encoding needs an 8-bit format; got {other}. Convert first."
130 )));
131 }
132 }
133 if desc.width == 0
134 || desc.height == 0
135 || desc.width > MAX_DIMENSION
136 || desc.height > MAX_DIMENSION
137 {
138 return Err(PixelsError::unsupported(format!(
139 "AVIF dimensions must be 1 to {MAX_DIMENSION}; {}x{} is not",
140 desc.width, desc.height
141 )));
142 }
143 let capacity = desc
144 .byte_len()
145 .ok_or_else(|| PixelsError::malformed("avif", "image size overflows"))?;
146 self.state = Some(State {
147 descriptor: *desc,
148 pixels: Vec::with_capacity(capacity),
149 rows_written: 0,
150 });
151 Ok(())
152 }
153
154 fn write_row(&mut self, row: &[u8], _sink: &mut dyn Sink) -> Result<()> {
155 let Some(state) = self.state.as_mut() else {
156 return Err(PixelsError::invalid_argument(
157 "row",
158 "write_row called before write_header",
159 ));
160 };
161 let expected = state.descriptor.row_bytes();
162 if row.len() != expected {
163 return Err(PixelsError::invalid_argument(
164 "row",
165 format!("row is {} bytes, expected {expected}", row.len()),
166 ));
167 }
168 if state.rows_written >= state.descriptor.height {
169 return Err(PixelsError::invalid_argument(
170 "row",
171 format!("more than {} rows written", state.descriptor.height),
172 ));
173 }
174 state.pixels.extend_from_slice(row);
175 state.rows_written += 1;
176 Ok(())
177 }
178
179 fn finish(&mut self, sink: &mut dyn Sink) -> Result<()> {
180 let Some(state) = self.state.as_ref() else {
181 return Err(PixelsError::invalid_argument(
182 "sink",
183 "finish called before write_header",
184 ));
185 };
186 if state.rows_written < state.descriptor.height {
187 return Err(PixelsError::malformed(
188 "avif",
189 format!(
190 "{} of {} rows were written",
191 state.rows_written, state.descriptor.height
192 ),
193 ));
194 }
195 let bytes = encode(
196 state,
197 &self.options,
198 self.icc.as_deref(),
199 self.exif.as_deref(),
200 self.xmp.as_deref(),
201 )?;
202 sink.write_all(&bytes)?;
203 sink.flush()
204 }
205}
206
207fn to_planes(
210 pixels: &[u8],
211 format: PixelFormat,
212 width: usize,
213 height: usize,
214 subsampled: bool,
215) -> (Vec<Vec<u16>>, Option<Vec<u16>>) {
216 let channels = format.channels();
217 let grey = matches!(format, PixelFormat::Gray8 | PixelFormat::GrayA8);
218 let has_alpha = matches!(format, PixelFormat::GrayA8 | PixelFormat::Rgba8);
219 let pixel = |i: usize| pixels.get(i * channels..(i + 1) * channels).unwrap_or(&[]);
220 let alpha = has_alpha.then(|| {
221 (0..width * height)
222 .map(|i| u16::from(pixel(i).get(channels - 1).copied().unwrap_or(255)))
223 .collect::<Vec<u16>>()
224 });
225 if grey {
226 let y = (0..width * height)
227 .map(|i| u16::from(pixel(i).first().copied().unwrap_or(0)))
228 .collect();
229 return (vec![y], alpha);
230 }
231 let mut y = Vec::with_capacity(width * height);
233 let mut u_full = Vec::with_capacity(width * height);
234 let mut v_full = Vec::with_capacity(width * height);
235 for i in 0..width * height {
236 let p = pixel(i);
237 let [r, g, b] = [0, 1, 2].map(|k| f64::from(p.get(k).copied().unwrap_or(0)));
238 let luma = 0.299 * r + 0.587 * g + 0.114 * b;
239 y.push(luma.round().clamp(0.0, 255.0) as u16);
240 u_full.push((b - luma) / 1.772);
241 v_full.push((r - luma) / 1.402);
242 }
243 if !subsampled {
244 let full = |plane: &[f64]| -> Vec<u16> {
245 plane
246 .iter()
247 .map(|&c| (c + 128.0).round().clamp(0.0, 255.0) as u16)
248 .collect()
249 };
250 return (vec![y, full(&u_full), full(&v_full)], alpha);
251 }
252 let (cw, ch) = (width.div_ceil(2), height.div_ceil(2));
253 let average = |full: &[f64]| -> Vec<u16> {
254 let mut out = Vec::with_capacity(cw * ch);
255 for cy in 0..ch {
256 for cx in 0..cw {
257 let (mut sum, mut n) = (0.0, 0.0);
258 for (dx, dy) in [(0, 0), (1, 0), (0, 1), (1, 1)] {
259 let (x, yy) = (cx * 2 + dx, cy * 2 + dy);
260 if x < width && yy < height {
261 sum += full.get(yy * width + x).copied().unwrap_or(0.0);
262 n += 1.0;
263 }
264 }
265 out.push((sum / n + 128.0).round().clamp(0.0, 255.0) as u16);
266 }
267 }
268 out
269 };
270 (vec![y, average(&u_full), average(&v_full)], alpha)
271}
272
273fn encode(
274 state: &State,
275 options: &EncodeOptions,
276 icc: Option<&[u8]>,
277 exif: Option<&[u8]>,
278 xmp: Option<&[u8]>,
279) -> Result<Vec<u8>> {
280 let (width, height) = (state.descriptor.width, state.descriptor.height);
281 let subsampled = options.chroma_subsampling != Some(ChromaSubsampling::Yuv444);
282 let effort = options.effort.unwrap_or(DEFAULT_EFFORT).min(9);
283 let (planes, alpha) = to_planes(
284 &state.pixels,
285 state.descriptor.pixel,
286 width as usize,
287 height as usize,
288 subsampled,
289 );
290 let qindex = qindex_for_quality(options.quality);
291 let refs: Vec<&[u16]> = planes.iter().map(Vec::as_slice).collect();
292 let colour = encode_still_with(
293 &Picture {
294 width,
295 height,
296 subsampled,
297 planes: &refs,
298 cicp: CICP,
299 full_range: true,
300 },
301 qindex,
302 effort,
303 )?;
304 let alpha = match alpha.filter(|a| a.iter().any(|&v| v != 255)) {
306 Some(alpha) => Some(encode_still_with(
307 &Picture {
308 width,
309 height,
310 subsampled: true,
311 planes: &[&alpha],
312 cicp: (2, 2, 2),
313 full_range: true,
314 },
315 (qindex / 2).max(1),
316 effort,
317 )?),
318 None => None,
319 };
320 Ok(container(
321 width,
322 height,
323 planes.len() == 1,
324 &colour,
325 alpha.as_ref(),
326 icc,
327 exif,
328 xmp,
329 ))
330}
331
332fn bx(kind: &[u8; 4], payload: &[u8]) -> Vec<u8> {
334 let mut out = Vec::with_capacity(payload.len() + 8);
335 out.extend_from_slice(&(payload.len() as u32 + 8).to_be_bytes());
336 out.extend_from_slice(kind);
337 out.extend_from_slice(payload);
338 out
339}
340
341fn full(kind: &[u8; 4], version: u8, flags: u32, payload: &[u8]) -> Vec<u8> {
343 let mut body = vec![version];
344 body.extend_from_slice(&flags.to_be_bytes()[1..]);
345 body.extend_from_slice(payload);
346 bx(kind, &body)
347}
348
349fn av1c(coded: &CodedStill, mono: bool) -> Vec<u8> {
351 let mut payload = vec![
352 0x81, (coded.profile << 5) | (coded.level & 0x1F),
354 (u8::from(mono) << 4) | if coded.subsampled { 0b1100 } else { 0 },
357 0, ];
359 payload.extend_from_slice(&coded.sequence_header_obu);
360 bx(b"av1C", &payload)
361}
362
363#[allow(
365 clippy::too_many_arguments,
366 reason = "the image, its alpha and each kind of metadata"
367)]
368fn container(
369 width: u32,
370 height: u32,
371 mono: bool,
372 colour: &CodedStill,
373 alpha: Option<&CodedStill>,
374 icc: Option<&[u8]>,
375 exif: Option<&[u8]>,
376 xmp: Option<&[u8]>,
377) -> Vec<u8> {
378 let ftyp = bx(b"ftyp", b"avif\0\0\0\0avifmif1miaf");
379 let ispe = {
380 let mut p = width.to_be_bytes().to_vec();
381 p.extend_from_slice(&height.to_be_bytes());
382 full(b"ispe", 0, 0, &p)
383 };
384 let pixi = |channels: u8| {
385 let mut p = vec![channels];
386 p.extend(std::iter::repeat_n(8, usize::from(channels)));
387 full(b"pixi", 0, 0, &p)
388 };
389 let colr = {
390 let mut p = b"nclx".to_vec();
391 for v in [CICP.0, CICP.1, CICP.2] {
392 p.extend_from_slice(&u16::from(v).to_be_bytes());
393 }
394 p.push(0x80); bx(b"colr", &p)
396 };
397 let mut ipco = [
400 ispe,
401 pixi(if mono { 1 } else { 3 }),
402 av1c(colour, mono),
403 colr,
404 ]
405 .concat();
406 let mut associations: Vec<(u16, Vec<u8>)> = vec![(1, vec![1, 2, 0x80 | 3, 4])];
407 if let Some(alpha) = alpha {
408 let mut urn = crate::meta::URN_ALPHA.as_bytes().to_vec();
409 urn.push(0);
410 ipco.extend_from_slice(&[pixi(1), av1c(alpha, true), full(b"auxC", 0, 0, &urn)].concat());
411 associations.push((2, vec![1, 5, 0x80 | 6, 7]));
412 }
413 if let Some(profile) = icc {
416 ipco.extend_from_slice(&bx(b"colr", &[b"prof".as_slice(), profile].concat()));
417 let index = if alpha.is_some() { 8 } else { 5 };
418 if let Some((_, props)) = associations.first_mut() {
419 props.push(index);
420 }
421 }
422 let ipma = {
423 let mut p = (associations.len() as u32).to_be_bytes().to_vec();
424 for (item, props) in &associations {
425 p.extend_from_slice(&item.to_be_bytes());
426 p.push(props.len() as u8);
427 p.extend_from_slice(props);
428 }
429 full(b"ipma", 0, 0, &p)
430 };
431 let iprp = bx(b"iprp", &[bx(b"ipco", &ipco), ipma].concat());
432
433 let hdlr = full(b"hdlr", 0, 0, b"\0\0\0\0pict\0\0\0\0\0\0\0\0\0\0\0\0\0");
434 let pitm = full(b"pitm", 0, 0, &1_u16.to_be_bytes());
435 let infe = |item: u16, kind: &[u8; 4]| {
436 let mut p = item.to_be_bytes().to_vec();
437 p.extend_from_slice(&[0, 0]);
438 p.extend_from_slice(kind);
439 p.push(0); if kind == b"mime" {
441 p.extend_from_slice(b"application/rdf+xml\0");
442 }
443 full(b"infe", 2, 0, &p)
444 };
445 let exif_data = exif.map(|block| [&[0_u8; 4][..], block].concat());
449 let mut items: Vec<(&[u8; 4], &[u8])> = std::iter::once(colour)
450 .chain(alpha)
451 .map(|coded| (b"av01", coded.data.as_slice()))
452 .collect();
453 if let Some(data) = &exif_data {
454 items.push((b"Exif", data));
455 }
456 if let Some(data) = xmp {
458 items.push((b"mime", data));
459 }
460 let iinf = {
461 let mut p = (items.len() as u16).to_be_bytes().to_vec();
462 for (i, (kind, _)) in items.iter().enumerate() {
463 p.extend_from_slice(&infe(i as u16 + 1, kind));
464 }
465 full(b"iinf", 0, 0, &p)
466 };
467 let mut references = Vec::new();
470 if alpha.is_some() {
471 references.extend(bx(b"auxl", &[0, 2, 0, 1, 0, 1]));
472 }
473 for (i, (kind, _)) in items.iter().enumerate() {
474 if matches!(*kind, b"Exif" | b"mime") {
475 let id = i as u16 + 1;
476 references.extend(bx(
477 b"cdsc",
478 &[&id.to_be_bytes()[..], &[0, 1, 0, 1]].concat(),
479 ));
480 }
481 }
482 let iref = (!references.is_empty()).then(|| full(b"iref", 0, 0, &references));
483 let iloc = |mdat_start: u32| {
485 let mut p = vec![0x44, 0x00];
486 p.extend_from_slice(&(items.len() as u16).to_be_bytes());
487 let mut offset = mdat_start;
488 for (i, (_, data)) in items.iter().enumerate() {
489 p.extend_from_slice(&(i as u16 + 1).to_be_bytes());
490 p.extend_from_slice(&[0, 0, 0, 1]); p.extend_from_slice(&offset.to_be_bytes());
492 p.extend_from_slice(&(data.len() as u32).to_be_bytes());
493 offset += data.len() as u32;
494 }
495 full(b"iloc", 0, 0, &p)
496 };
497 let meta = |mdat_start: u32| {
498 let mut p = [hdlr.clone(), pitm.clone(), iloc(mdat_start), iinf.clone()].concat();
499 if let Some(iref) = &iref {
500 p.extend_from_slice(iref);
501 }
502 p.extend_from_slice(&iprp);
503 full(b"meta", 0, 0, &p)
504 };
505 let mdat_start = (ftyp.len() + meta(0).len() + 8) as u32;
507 let mdat_payload: Vec<u8> = items
508 .iter()
509 .flat_map(|(_, data)| data.iter().copied())
510 .collect();
511 [ftyp, meta(mdat_start), bx(b"mdat", &mdat_payload)].concat()
512}
513
514#[cfg(test)]
515#[allow(
516 clippy::unwrap_used,
517 clippy::indexing_slicing,
518 reason = "tests operate on known-good values"
519)]
520mod tests {
521 use super::*;
522
523 #[test]
524 fn rgb_to_yuv_follows_bt601_full_range() {
525 let pixels = [255, 255, 255, 0, 0, 0, 255, 0, 0, 0, 0, 255];
526 let (planes, alpha) = to_planes(&pixels, PixelFormat::Rgb8, 2, 2, true);
527 assert!(alpha.is_none());
528 assert_eq!(planes[0], vec![255, 0, 76, 29]);
529 let u = (((-76.245_f64) + 255.0 * (1.0 - 0.114)) / 1.772 / 4.0 + 128.0).round();
531 assert!(
532 (i32::from(planes[1][0]) - u as i32).abs() <= 1,
533 "{} vs {u}",
534 planes[1][0]
535 );
536 assert_eq!(planes[1].len(), 1);
537 }
538
539 #[test]
540 fn grey_keeps_one_plane_and_alpha_is_split_out() {
541 let pixels = [10, 200, 20, 255, 30, 0];
542 let (planes, alpha) = to_planes(&pixels, PixelFormat::GrayA8, 3, 1, true);
543 assert_eq!(planes, vec![vec![10, 20, 30]]);
544 assert_eq!(alpha.unwrap(), vec![200, 255, 0]);
545 }
546
547 #[test]
548 fn odd_sizes_average_only_the_samples_present() {
549 let pixels = [0, 0, 255, 0, 0, 255, 0, 0, 255];
550 let (planes, _) = to_planes(&pixels, PixelFormat::Rgb8, 3, 1, true);
551 assert_eq!((planes[1].len(), planes[2].len()), (2, 2));
552 assert_eq!(planes[1][0], planes[1][1]);
553 }
554
555 fn encode_with(descriptor: &ImageDescriptor, pixels: &[u8], options: EncodeOptions) -> Vec<u8> {
557 let mut encoder = AvifEncoder::from_options(&options);
558 let mut out = Vec::new();
559 encoder.write_header(descriptor, &mut out).unwrap();
560 for row in pixels.chunks_exact(descriptor.row_bytes()) {
561 encoder.write_row(row, &mut out).unwrap();
562 }
563 encoder.finish(&mut out).unwrap();
564 out
565 }
566
567 fn decode(bytes: &[u8]) -> Vec<u8> {
568 use otf_pixels_core::{Decoder, Limits};
569 let mut decoder = crate::AvifDecoder::new(bytes, Limits::default()).unwrap();
570 let mut row = vec![0_u8; decoder.descriptor().row_bytes()];
571 let mut out = Vec::new();
572 for _ in 0..decoder.descriptor().height {
573 decoder.read_row(&mut row).unwrap();
574 out.extend_from_slice(&row);
575 }
576 out
577 }
578
579 fn av1c_bytes(file: &[u8]) -> (u8, u8) {
581 let at = file.windows(4).position(|w| w == b"av1C").unwrap() + 4;
582 (file[at + 1] >> 5, file[at + 2] & 0b1100)
583 }
584
585 #[test]
586 fn full_resolution_chroma_keeps_every_chroma_sample() {
587 let pixels = [255, 255, 255, 0, 0, 0, 255, 0, 0, 0, 0, 255];
588 let (planes, _) = to_planes(&pixels, PixelFormat::Rgb8, 2, 2, false);
589 assert_eq!((planes[1].len(), planes[2].len()), (4, 4));
590 assert_eq!(planes[1][0], 128, "white has no chroma");
591 assert_eq!(planes[2][2], 255, "red's V is full scale");
592 }
593
594 #[test]
595 fn four_four_four_is_profile_one_and_sharper_in_colour() {
596 let descriptor = ImageDescriptor::new(32, 16, PixelFormat::Rgb8).unwrap();
598 let pixels: Vec<u8> = (0..32 * 16)
599 .flat_map(|i| if i % 2 == 0 { [255, 0, 0] } else { [0, 0, 255] })
600 .collect();
601 let options = EncodeOptions::with_quality(90).unwrap();
602 let sub = encode_with(&descriptor, &pixels, options);
603 let full = encode_with(
604 &descriptor,
605 &pixels,
606 options.with_chroma_subsampling(ChromaSubsampling::Yuv444),
607 );
608 assert_eq!(av1c_bytes(&sub), (0, 0b1100));
609 assert_eq!(av1c_bytes(&full), (1, 0));
610 assert_eq!(
611 encode_with(
612 &descriptor,
613 &pixels,
614 options.with_chroma_subsampling(ChromaSubsampling::Yuv420)
615 ),
616 sub
617 );
618 let error = |bytes: &[u8]| -> u64 {
619 decode(bytes)
620 .iter()
621 .zip(&pixels)
622 .map(|(a, b)| u64::from(a.abs_diff(*b)))
623 .sum()
624 };
625 assert!(
626 error(&full) * 2 < error(&sub),
627 "{} against {}",
628 error(&full),
629 error(&sub)
630 );
631 }
632
633 #[test]
634 fn grey_and_alpha_stay_profile_zero_under_four_four_four() {
635 let descriptor = ImageDescriptor::new(9, 7, PixelFormat::GrayA8).unwrap();
636 let pixels: Vec<u8> = (0..9 * 7)
637 .flat_map(|i| [(i * 4) as u8, (i * 3) as u8])
638 .collect();
639 let file = encode_with(
640 &descriptor,
641 &pixels,
642 EncodeOptions::default().with_chroma_subsampling(ChromaSubsampling::Yuv444),
643 );
644 assert_eq!(av1c_bytes(&file).0, 0);
645 assert_eq!(decode(&file).len(), 9 * 7 * 2);
646 }
647
648 #[test]
649 fn effort_narrows_the_search_and_four_is_the_default() {
650 let descriptor = ImageDescriptor::new(40, 24, PixelFormat::Rgb8).unwrap();
651 let pixels: Vec<u8> = (0..40 * 24 * 3)
652 .map(|i| ((i * 13 + i / 97) % 256) as u8)
653 .collect();
654 let base = EncodeOptions::default();
655 let default = encode_with(&descriptor, &pixels, base);
656 for effort in 3..=10 {
657 assert_eq!(
658 encode_with(&descriptor, &pixels, base.with_effort(effort).unwrap()),
659 default,
660 "effort {effort}"
661 );
662 }
663 let fast = encode_with(&descriptor, &pixels, base.with_effort(0).unwrap());
664 assert_ne!(fast, default);
665 assert_eq!(decode(&fast).len(), pixels.len());
666 }
667}