aom_decode/
avif.rs

1//! Wrapper for decoding AV1 data in AVIF
2use crate::color::{ChromaSampling, Depth, MatrixCoefficients, Range};
3use crate::{Config, Decoder, FrameTempRef, Planes};
4use crate::{Error, Result};
5use imgref::{Img, ImgRef, ImgVec};
6use rgb::prelude::*;
7use rgb::{Rgb, Rgba};
8use std::borrow::Cow;
9use yuv::{YuvPlanarImage, YuvRange, YuvStandardMatrix};
10
11/// AVIF decoder + converter
12pub struct Avif {
13    decoder: Decoder,
14    avif: ParsedAvifData,
15}
16
17/// A decoded image
18pub enum Image {
19    RGB8(ImgVec<Rgb<u8>>),
20    RGBA8(ImgVec<Rgba<u8>>),
21    RGB16(ImgVec<Rgb<u16>>),
22    RGBA16(ImgVec<Rgba<u16>>),
23    Gray8(ImgVec<u8>),
24    Gray16(ImgVec<u16>),
25}
26
27/// Raw data (you can inspect the metadata)
28pub use avif_parse::AvifData as ParsedAvifData;
29/// See [`ParsedAvifData::primary_item_metadata()`]
30///
31pub use avif_parse::AV1Metadata;
32pub use avif_parse::Error as AvifParseError;
33pub use yuv::YuvError;
34
35impl Avif {
36    pub fn decode(data: &[u8], config: &Config) -> Result<Self> {
37        Self::from_parsed_avif_data(Self::parse_avif(data)?, config)
38    }
39
40    pub fn convert(&mut self) -> Result<Image> {
41        // aom decoder recycles buffers, so can't have both color and alpha without copying,
42        // therefore conversion will put placeholders and then update alpha
43        let has_alpha = self.avif.alpha_item.is_some();
44        let color = self.raw_color_data()?;
45        let mut img = match color.planes()? {
46            Planes::YuvPlanes8 { y, u, v, chroma_sampling } => {
47                yuv_to_rgb8(&color, y, chroma_sampling, u, v, has_alpha)?
48            },
49            Planes::Mono8(y) => yuv_to_gray8(&color, y, has_alpha),
50            Planes::Mono16(y, depth) => yuv_to_gray16(&color, depth, y, has_alpha),
51            Planes::YuvPlanes16 { y, u, v, chroma_sampling, depth } => {
52                yuv_to_rgb16(&color, depth, y, chroma_sampling, u, v, has_alpha)?
53            },
54        };
55        let color_mc = color.matrix_coefficients().unwrap_or(MatrixCoefficients::Identity);
56
57        let premultiplied_alpha = self.avif.premultiplied_alpha;
58        if let Some(alpha) = self.raw_alpha_data()? {
59            let range = alpha.range();
60            if let Some(alpha_mc) = alpha.matrix_coefficients().filter(|&mc| mc != MatrixCoefficients::Identity)
61                && color_mc != alpha_mc {
62                return Err(Error::Unsupported("alpha image has color info"));
63            }
64            match alpha.planes()? {
65                Planes::YuvPlanes8 { y, .. } | Planes::Mono8(y) => {
66                    add_alpha8(&mut img, y, range, premultiplied_alpha)?;
67                },
68                Planes::YuvPlanes16 { y, depth, .. } | Planes::Mono16(y, depth) => {
69                    add_alpha16(&mut img, y, depth, range, premultiplied_alpha)?;
70                },
71            }
72        } else if has_alpha {
73            return Err(Error::Unsupported("invalid alpha"));
74        }
75        Ok(img)
76    }
77
78    pub fn parse_avif(data: &[u8]) -> Result<ParsedAvifData> {
79        Ok(avif_parse::read_avif(&mut &data[..])?)
80    }
81
82    pub fn from_parsed_avif_data(avif: ParsedAvifData, config: &Config) -> Result<Self> {
83        let decoder = Decoder::new(config)?;
84        Ok(Self { decoder, avif })
85    }
86
87    pub fn raw_color_data(&mut self) -> Result<FrameTempRef<'_>> {
88        self.decoder.decode_frame(&self.avif.primary_item)
89    }
90
91    pub fn raw_alpha_data(&mut self) -> Result<Option<FrameTempRef<'_>>> {
92        Ok(if let Some(alpha) = &self.avif.alpha_item {
93            Some(self.decoder.decode_frame(alpha)?)
94        } else {
95            None
96        })
97    }
98}
99
100fn add_alpha16(img: &mut Image, y: ImgRef<'_, [u8; 2]>, depth: Depth, range: Range, premultiplied_alpha: bool) -> Result<()> {
101    match img {
102        Image::RGBA8(img) => {
103            for (y_row, img_row) in y.rows().zip(img.rows_mut()) {
104                if y_row.len() != img_row.len() {
105                    return Err(Error::Unsupported("invalid alpha size"));
106                }
107                for (y, px) in y_row.iter().copied().zip(img_row.iter_mut()) {
108                    px.a = (luma16(u16::from_ne_bytes(y), depth, range) >> 8) as u8;
109                }
110                if premultiplied_alpha {
111                    unpremultiply8(img_row);
112                }
113            }
114        },
115        Image::RGBA16(img) => {
116            for (y_row, img_row) in y.rows().zip(img.rows_mut()) {
117                if y_row.len() != img_row.len() {
118                    return Err(Error::Unsupported("invalid alpha size"));
119                }
120                for (y, px) in y_row.iter().copied().zip(img_row.iter_mut()) {
121                    px.a = luma16(u16::from_ne_bytes(y), depth, range);
122                }
123                if premultiplied_alpha {
124                    unpremultiply16(img_row);
125                }
126            }
127        },
128        _ => return Err(Error::Unsupported("internal error")),
129    }
130    Ok(())
131}
132
133fn add_alpha8(img: &mut Image, y: ImgRef<'_, u8>, range: Range, premultiplied_alpha: bool) -> Result<()> {
134    match img {
135        Image::RGBA8(img) => {
136            for (y_row, img_row) in y.rows().zip(img.rows_mut()) {
137                if y_row.len() != img_row.len() {
138                    return Err(Error::Unsupported("invalid alpha size"));
139                }
140                for (y, px) in y_row.iter().copied().zip(img_row.iter_mut()) {
141                    px.a = luma8(y, range);
142                }
143                if premultiplied_alpha {
144                    unpremultiply8(img_row);
145                }
146            }
147        },
148        Image::RGBA16(img) => {
149            for (y_row, img_row) in y.rows().zip(img.rows_mut()) {
150                if y_row.len() != img_row.len() {
151                    return Err(Error::Unsupported("invalid alpha size"));
152                }
153                for (y, px) in y_row.iter().copied().zip(img_row.iter_mut()) {
154                    px.a = luma16(u16::from(y), Depth::Depth8, range);
155                }
156                if premultiplied_alpha {
157                    unpremultiply16(img_row);
158                }
159            }
160        },
161        _ => return Err(Error::Unsupported("internal error")),
162    }
163    Ok(())
164}
165
166#[inline(never)]
167fn unpremultiply8(img_row: &mut [Rgba<u8>]) {
168    for px in img_row.iter_mut() {
169        if px.a != 255 && px.a != 0 {
170            *px.rgb_mut() = px.rgb().map(|c| (u16::from(c) * 255 / u16::from(px.a)).min(255) as u8);
171        }
172    }
173}
174
175#[inline(never)]
176fn unpremultiply16(img_row: &mut [Rgba<u16>]) {
177    for px in img_row.iter_mut() {
178        if px.a != 0xFFFF && px.a != 0 {
179            *px.rgb_mut() = px.rgb().map(|c| (u32::from(c) * 0xFFFF / u32::from(px.a)).min(0xFFFF) as u16);
180        }
181    }
182}
183
184fn imgref_align_to_u16(plane: ImgRef<'_, [u8; 2]>) -> Img<Cow<'_, [u16]>> {
185    plane.new_buf(match bytemuck::try_cast_slice(plane.buf()) {
186        Ok(samples) => Cow::Borrowed(samples),
187        Err(_) => Cow::Owned(plane.buf().iter().map(|b| u16::from_ne_bytes(*b)).collect()),
188    })
189}
190
191fn yuv_to_rgb16(color: &FrameTempRef, depth: Depth, y: ImgRef<'_, [u8; 2]>, chroma_sampling: ChromaSampling, u: ImgRef<'_, [u8; 2]>, v: ImgRef<'_, [u8; 2]>, has_alpha: bool) -> Result<Image, Error> {
192    let mc = color.matrix_coefficients().unwrap_or(MatrixCoefficients::BT601);
193    let conv = conversion(mc, color.range())?;
194    let range = to_yuv_range(color.range());
195    let width = y.width();
196    let height = y.height();
197    // Planes are read where the decoder left them; only unaligned planes are copied
198    let (y_plane, u_plane, v_plane) = (imgref_align_to_u16(y), imgref_align_to_u16(u), imgref_align_to_u16(v));
199    let planar = YuvPlanarImage {
200        y_plane: y_plane.buf(),
201        y_stride: y_plane.stride() as u32,
202        u_plane: u_plane.buf(),
203        u_stride: u_plane.stride() as u32,
204        v_plane: v_plane.buf(),
205        v_stride: v_plane.stride() as u32,
206        width: width as u32,
207        height: height as u32,
208    };
209    if has_alpha {
210        let mut out = vec![Rgba::<u16>::new(0, 0, 0, 0); width * height];
211        convert16(&planar, bytemuck::cast_slice_mut(out.as_mut_slice()), (width * 4) as u32, chroma_sampling, depth, conv, range, true)?;
212        Ok(Image::RGBA16(ImgVec::new(out, width, height)))
213    } else {
214        let mut out = vec![Rgb::<u16>::new(0, 0, 0); width * height];
215        convert16(&planar, bytemuck::cast_slice_mut(out.as_mut_slice()), (width * 3) as u32, chroma_sampling, depth, conv, range, false)?;
216        Ok(Image::RGB16(ImgVec::new(out, width, height)))
217    }
218}
219
220fn yuv_to_gray16(color: &FrameTempRef, depth: Depth, y: ImgRef<'_, [u8; 2]>, has_alpha: bool) -> Image {
221    let range = color.range();
222    let width = y.width();
223    let height = y.height();
224    if has_alpha {
225        let mut out = Vec::with_capacity(width * height);
226        out.extend(y.rows().flat_map(|row| {
227            row.iter().copied().map(|y| {
228                let g = luma16(u16::from_ne_bytes(y), depth, range);
229                Rgba::new(g, g, g, 0)
230            })
231        }));
232        Image::RGBA16(ImgVec::new(out, width, height))
233    } else {
234        let mut out = Vec::with_capacity(width * height);
235        out.extend(y.rows().flat_map(|row| {
236            row.iter()
237                .copied()
238                .map(|y| luma16(u16::from_ne_bytes(y), depth, range))
239        }));
240        Image::Gray16(ImgVec::new(out, width, height))
241    }
242}
243
244fn yuv_to_gray8(color: &FrameTempRef, y: ImgRef<'_, u8>, has_alpha: bool) -> Image {
245    let range = color.range();
246    let width = y.width();
247    let height = y.height();
248    if has_alpha {
249        let mut out = Vec::with_capacity(width * height);
250        out.extend(y.rows().flat_map(|row| {
251            row.iter().copied().map(|y| {
252                let g = luma8(y, range);
253                Rgba::new(g, g, g, 0)
254            })
255        }));
256        Image::RGBA8(ImgVec::new(out, width, height))
257    } else {
258        let mut out = Vec::with_capacity(width * height);
259        out.extend(y.rows().flat_map(|row| {
260            row.iter()
261                .copied()
262                .map(|y| luma8(y, range))
263        }));
264        Image::Gray8(ImgVec::new(out, width, height))
265    }
266}
267
268fn yuv_to_rgb8(color: &FrameTempRef, y: ImgRef<'_, u8>, chroma_sampling: ChromaSampling, u: ImgRef<'_, u8>, v: ImgRef<'_, u8>, has_alpha: bool) -> Result<Image, Error> {
269    let mc = color.matrix_coefficients().unwrap_or(MatrixCoefficients::BT601);
270    let conv = conversion(mc, color.range())?;
271    let range = to_yuv_range(color.range());
272    let width = y.width();
273    let height = y.height();
274    let planar = YuvPlanarImage {
275        y_plane: y.buf(),
276        y_stride: y.stride() as u32,
277        u_plane: u.buf(),
278        u_stride: u.stride() as u32,
279        v_plane: v.buf(),
280        v_stride: v.stride() as u32,
281        width: width as u32,
282        height: height as u32,
283    };
284    if has_alpha {
285        let mut out = vec![Rgba::<u8>::new(0, 0, 0, 0); width * height];
286        convert8(&planar, bytemuck::cast_slice_mut(out.as_mut_slice()), (width * 4) as u32, chroma_sampling, conv, range, true)?;
287        Ok(Image::RGBA8(ImgVec::new(out, width, height)))
288    } else {
289        let mut out = vec![Rgb::<u8>::new(0, 0, 0); width * height];
290        convert8(&planar, bytemuck::cast_slice_mut(out.as_mut_slice()), (width * 3) as u32, chroma_sampling, conv, range, false)?;
291        Ok(Image::RGB8(ImgVec::new(out, width, height)))
292    }
293}
294
295/// Meaning of the frame's planes, as far as conversion to RGB is concerned
296#[derive(Debug, Copy, Clone, PartialEq)]
297enum Conversion {
298    /// Planes are G, B, R (identity matrix). Only valid without chroma subsampling.
299    Gbr,
300    /// `YCgCo`
301    YCgCo,
302    /// YCbCr with one of the standard matrices
303    Matrix(YuvStandardMatrix),
304}
305
306fn conversion(mc: MatrixCoefficients, range: Range) -> Result<Conversion, Error> {
307    Ok(match mc {
308        MatrixCoefficients::Identity => Conversion::Gbr,
309        MatrixCoefficients::YCgCo => Conversion::YCgCo,
310        MatrixCoefficients::BT709 => Conversion::Matrix(YuvStandardMatrix::Bt709),
311        MatrixCoefficients::FCC => Conversion::Matrix(YuvStandardMatrix::Fcc),
312        // BT.470 System B,G uses the same matrix as BT.601 (ITU-T H.273 table 4),
313        // which is not what `yuv`'s own `Bt470_6` variant means
314        MatrixCoefficients::BT470BG |
315        MatrixCoefficients::BT601 => Conversion::Matrix(YuvStandardMatrix::Bt601),
316        MatrixCoefficients::SMPTE240 => Conversion::Matrix(YuvStandardMatrix::Smpte240),
317        MatrixCoefficients::BT2020NCL => Conversion::Matrix(YuvStandardMatrix::Bt2020),
318        _ => {
319            log::debug!("Unsupported matrix coefficients: {mc:?} ({range:?})");
320            return Err(Error::Unsupported("matrix coefficients"));
321        },
322    })
323}
324
325#[inline]
326fn to_yuv_range(range: Range) -> YuvRange {
327    match range {
328        Range::Full => YuvRange::Full,
329        Range::Limited => YuvRange::Limited,
330    }
331}
332
333/// Convert 8-bit planes into `dst`, which is RGB when `has_alpha` is false and RGBA otherwise
334fn convert8(planar: &YuvPlanarImage<u8>, dst: &mut [u8], dst_stride: u32, chroma_sampling: ChromaSampling, conv: Conversion, range: YuvRange, has_alpha: bool) -> Result<(), Error> {
335    use ChromaSampling::{Cs420, Cs422, Cs444, Monochrome};
336    use Conversion::{Gbr, Matrix, YCgCo};
337    let res = match (conv, chroma_sampling) {
338        (Gbr, Cs444) => if has_alpha {
339            yuv::gbr_to_rgba(planar, dst, dst_stride, range)
340        } else {
341            yuv::gbr_to_rgb(planar, dst, dst_stride, range)
342        },
343        (Gbr, _) => return Err(Error::Unsupported("identity matrix with chroma subsampling")),
344        (YCgCo, Cs444) => if has_alpha {
345            yuv::ycgco444_to_rgba(planar, dst, dst_stride, range)
346        } else {
347            yuv::ycgco444_to_rgb(planar, dst, dst_stride, range)
348        },
349        (YCgCo, Cs422) => if has_alpha {
350            yuv::ycgco422_to_rgba(planar, dst, dst_stride, range)
351        } else {
352            yuv::ycgco422_to_rgb(planar, dst, dst_stride, range)
353        },
354        (YCgCo, Cs420) => if has_alpha {
355            yuv::ycgco420_to_rgba(planar, dst, dst_stride, range)
356        } else {
357            yuv::ycgco420_to_rgb(planar, dst, dst_stride, range)
358        },
359        (YCgCo, Monochrome) => unreachable!(),
360        (Matrix(m), Cs444) => if has_alpha {
361            yuv::yuv444_to_rgba(planar, dst, dst_stride, range, m)
362        } else {
363            yuv::yuv444_to_rgb(planar, dst, dst_stride, range, m)
364        },
365        (Matrix(m), Cs422) => if has_alpha {
366            yuv::yuv422_to_rgba(planar, dst, dst_stride, range, m)
367        } else {
368            yuv::yuv422_to_rgb(planar, dst, dst_stride, range, m)
369        },
370        (Matrix(m), Cs420) => if has_alpha {
371            yuv::yuv420_to_rgba(planar, dst, dst_stride, range, m)
372        } else {
373            yuv::yuv420_to_rgb(planar, dst, dst_stride, range, m)
374        },
375        (Matrix(_), Monochrome) => unreachable!(),
376    };
377    res.inspect_err(|e| log::debug!("{e} (conversion={conv:?}, range={range:?})"))?;
378    Ok(())
379}
380
381/// Convert 16-bit planes into `dst`, which is RGB when `has_alpha` is false and RGBA otherwise
382fn convert16(planar: &YuvPlanarImage<u16>, dst: &mut [u16], dst_stride: u32, chroma_sampling: ChromaSampling, depth: Depth, conv: Conversion, range: YuvRange, has_alpha: bool) -> Result<(), Error> {
383    use ChromaSampling::{Cs420, Cs422, Cs444, Monochrome};
384    use Conversion::{Gbr, Matrix, YCgCo};
385    let res = match (conv, chroma_sampling, depth) {
386        (Gbr, Cs444, Depth::Depth10) => if has_alpha {
387            yuv::gb10_to_rgba10(planar, dst, dst_stride, range)
388        } else {
389            yuv::gb10_to_rgb10(planar, dst, dst_stride, range)
390        },
391        (Gbr, Cs444, Depth::Depth12) => if has_alpha {
392            yuv::gb12_to_rgba12(planar, dst, dst_stride, range)
393        } else {
394            yuv::gb12_to_rgb12(planar, dst, dst_stride, range)
395        },
396        (Gbr, Cs444, _) => if has_alpha {
397            yuv::gb16_to_rgba16(planar, dst, dst_stride, range)
398        } else {
399            yuv::gb16_to_rgb16(planar, dst, dst_stride, range)
400        },
401        (Gbr, _, _) => return Err(Error::Unsupported("identity matrix with chroma subsampling")),
402        (YCgCo, Cs444, Depth::Depth10) => if has_alpha {
403            yuv::icgc410_to_rgba10(planar, dst, dst_stride, range)
404        } else {
405            yuv::icgc410_to_rgb10(planar, dst, dst_stride, range)
406        },
407        (YCgCo, Cs422, Depth::Depth10) => if has_alpha {
408            yuv::icgc210_to_rgba10(planar, dst, dst_stride, range)
409        } else {
410            yuv::icgc210_to_rgb10(planar, dst, dst_stride, range)
411        },
412        (YCgCo, Cs420, Depth::Depth10) => if has_alpha {
413            yuv::icgc010_to_rgba10(planar, dst, dst_stride, range)
414        } else {
415            yuv::icgc010_to_rgb10(planar, dst, dst_stride, range)
416        },
417        (YCgCo, Cs444, Depth::Depth12) => if has_alpha {
418            yuv::icgc412_to_rgba12(planar, dst, dst_stride, range)
419        } else {
420            yuv::icgc412_to_rgb12(planar, dst, dst_stride, range)
421        },
422        (YCgCo, Cs422, Depth::Depth12) => if has_alpha {
423            yuv::icgc212_to_rgba12(planar, dst, dst_stride, range)
424        } else {
425            yuv::icgc212_to_rgb12(planar, dst, dst_stride, range)
426        },
427        (YCgCo, Cs420, Depth::Depth12) => if has_alpha {
428            yuv::icgc012_to_rgba12(planar, dst, dst_stride, range)
429        } else {
430            yuv::icgc012_to_rgb12(planar, dst, dst_stride, range)
431        },
432        (YCgCo, _, _) => return Err(Error::Unsupported("YCgCo at this depth")),
433        (Matrix(m), Cs444, Depth::Depth10) => if has_alpha {
434            yuv::i410_to_rgba10(planar, dst, dst_stride, range, m)
435        } else {
436            yuv::i410_to_rgb10(planar, dst, dst_stride, range, m)
437        },
438        (Matrix(m), Cs444, Depth::Depth12) => if has_alpha {
439            yuv::i412_to_rgba12(planar, dst, dst_stride, range, m)
440        } else {
441            yuv::i412_to_rgb12(planar, dst, dst_stride, range, m)
442        },
443        (Matrix(m), Cs444, _) => if has_alpha {
444            yuv::i416_to_rgba16(planar, dst, dst_stride, range, m)
445        } else {
446            yuv::i416_to_rgb16(planar, dst, dst_stride, range, m)
447        },
448        (Matrix(m), Cs422, Depth::Depth10) => if has_alpha {
449            yuv::i210_to_rgba10(planar, dst, dst_stride, range, m)
450        } else {
451            yuv::i210_to_rgb10(planar, dst, dst_stride, range, m)
452        },
453        (Matrix(m), Cs422, Depth::Depth12) => if has_alpha {
454            yuv::i212_to_rgba12(planar, dst, dst_stride, range, m)
455        } else {
456            yuv::i212_to_rgb12(planar, dst, dst_stride, range, m)
457        },
458        (Matrix(m), Cs422, _) => if has_alpha {
459            yuv::i216_to_rgba16(planar, dst, dst_stride, range, m)
460        } else {
461            yuv::i216_to_rgb16(planar, dst, dst_stride, range, m)
462        },
463        (Matrix(m), Cs420, Depth::Depth10) => if has_alpha {
464            yuv::i010_to_rgba10(planar, dst, dst_stride, range, m)
465        } else {
466            yuv::i010_to_rgb10(planar, dst, dst_stride, range, m)
467        },
468        (Matrix(m), Cs420, Depth::Depth12) => if has_alpha {
469            yuv::i012_to_rgba12(planar, dst, dst_stride, range, m)
470        } else {
471            yuv::i012_to_rgb12(planar, dst, dst_stride, range, m)
472        },
473        (Matrix(m), Cs420, _) => if has_alpha {
474            yuv::i016_to_rgba16(planar, dst, dst_stride, range, m)
475        } else {
476            yuv::i016_to_rgb16(planar, dst, dst_stride, range, m)
477        },
478        (Matrix(_), Monochrome, _) => return Err(Error::Unsupported("unreachable")),
479    };
480    res.inspect_err(|e| log::debug!("{e} (conversion={conv:?}, range={range:?})"))?;
481
482    // The conversions work at their own bit depth, but the crate assumes full 16-bit range
483    if depth != Depth::Depth16 {
484        let shift = 16 - depth.bits();
485        let low_shift = depth.bits() - shift;
486        for v in dst.iter_mut() {
487            *v = (*v << shift) | (*v >> low_shift);
488        }
489    }
490    Ok(())
491}
492
493/// Expand an 8-bit luma sample to the full 8-bit range
494#[inline]
495fn luma8(v: u8, range: Range) -> u8 {
496    (luma16(u16::from(v), Depth::Depth8, range) >> 8) as u8
497}
498
499/// Expand a luma sample of the given depth to the full 16-bit range
500#[inline]
501fn luma16(v: u16, depth: Depth, range: Range) -> u16 {
502    let bits = depth.bits();
503    let (min, span) = match range {
504        Range::Full => (0, (1 << bits) - 1),
505        Range::Limited => (16 << (bits - 8), 219 << (bits - 8)),
506    };
507    let v = u32::from(v).saturating_sub(min).min(span);
508    ((v * 65535 + span / 2) / span) as u16
509}
aom_decode/avif.rs

aom_decode/
avif.rs
