aom_decode/
avif.rs

1//! Wrapper for decoding AV1 data in AVIF
2use crate::chroma::yuv_420;
3use crate::chroma::yuv_422;
4use crate::chroma::yuv_444;
5use crate::color::{ChromaSampling, MatrixCoefficients};
6use crate::Config;
7use crate::Decoder;
8use crate::Error;
9use crate::FrameTempRef;
10use crate::Result;
11use crate::RowsIter;
12use crate::RowsIters;
13use imgref::ImgVec;
14use yuv::color::{Depth, Range};
15use yuv::convert::RGBConvert;
16use yuv::RGB;
17use yuv::RGBA;
18use yuv::YUV;
19use rgb::prelude::*;
20
21/// AVIF decoder + converter
22pub struct Avif {
23    decoder: Decoder,
24    avif: ParsedAvifData,
25}
26
27/// A decoded image
28pub enum Image {
29    RGB8(ImgVec<RGB<u8>>),
30    RGBA8(ImgVec<RGBA<u8>>),
31    RGB16(ImgVec<RGB<u16>>),
32    RGBA16(ImgVec<RGBA<u16>>),
33    Gray8(ImgVec<u8>),
34    Gray16(ImgVec<u16>),
35}
36
37/// Raw data (you can inspect the metadata)
38pub use avif_parse::AvifData as ParsedAvifData;
39/// See [`ParsedAvifData::primary_item_metadata()`]
40///
41pub use avif_parse::AV1Metadata;
42pub use avif_parse::Error as AvifParseError;
43pub use yuv::Error as YuvError;
44
45impl Avif {
46    pub fn decode(data: &[u8], config: &Config) -> Result<Self> {
47        Self::from_parsed_avif_data(Self::parse_avif(data)?, config)
48    }
49
50    pub fn convert(&mut self) -> Result<Image> {
51        // aom decoder recycles buffers, so can't have both color and alpha without copying,
52        // therefore conversion will put placeholders and then update alpha
53        let has_alpha = self.avif.alpha_item.is_some();
54        let color = self.raw_color_data()?;
55        let range = color.range();
56        let mut img = match color.rows_iter()? {
57            RowsIters::YuvPlanes8 {y,u,v,chroma_sampling} => {
58                yuv_to_rgb8(&color, range, y, chroma_sampling, u, v, has_alpha)?
59            },
60            RowsIters::Mono8(y) => {
61                yuv_to_gray8(&color, range, y, has_alpha)?
62            },
63            RowsIters::Mono16(y, depth) => {
64                yuv_to_gray16(&color, range, depth, y, has_alpha)?
65            },
66            RowsIters::YuvPlanes16 {y,u,v,chroma_sampling, depth} => {
67                yuv_to_rgb16(&color, range, depth, y, chroma_sampling, u, v, has_alpha)?
68            },
69        };
70        let color_mc = color.matrix_coefficients().unwrap_or(MatrixCoefficients::Identity);
71        drop(color);
72
73        let premultiplied_alpha = self.avif.premultiplied_alpha;
74        if let Some(alpha) = self.raw_alpha_data()? {
75            let range = alpha.range();
76            if let Some(alpha_mc) = alpha.matrix_coefficients().filter(|&mc| mc != MatrixCoefficients::Identity) {
77                log::warn!("The alpha channel is not allowed to have a color matrix: {alpha_mc:?}");
78                if color_mc != alpha_mc {
79                    return Err(Error::Unsupported("alpha image has color info"));
80                }
81            }
82            let mc = MatrixCoefficients::Identity;
83            match alpha.rows_iter()? {
84                RowsIters::YuvPlanes8 { y, .. } | RowsIters::Mono8(y) => {
85                    let conv = RGBConvert::<u8>::new(range, mc)?;
86                    add_alpha8(&mut img, y, conv, premultiplied_alpha)?;
87                },
88                RowsIters::YuvPlanes16 { y, depth, .. } | RowsIters::Mono16(y, depth) => {
89                    let conv = RGBConvert::<u16>::new(range, mc, depth)?;
90                    add_alpha16(&mut img, y, conv, premultiplied_alpha)?;
91                },
92            };
93        } else if has_alpha {
94            return Err(Error::Unsupported("invalid alpha"));
95        }
96        Ok(img)
97    }
98
99    pub fn parse_avif(data: &[u8]) -> Result<ParsedAvifData> {
100        Ok(avif_parse::read_avif(&mut &data[..])?)
101    }
102
103    pub fn from_parsed_avif_data(avif: ParsedAvifData, config: &Config) -> Result<Self> {
104        let decoder = Decoder::new(config)?;
105        Ok(Self { decoder, avif })
106    }
107
108    pub fn raw_color_data(&mut self) -> Result<FrameTempRef<'_>> {
109        self.decoder.decode_frame(&self.avif.primary_item)
110    }
111
112    pub fn raw_alpha_data(&mut self) -> Result<Option<FrameTempRef<'_>>> {
113        Ok(if let Some(alpha) = &self.avif.alpha_item {
114            Some(self.decoder.decode_frame(alpha)?)
115        } else {
116            None
117        })
118    }
119}
120
121fn add_alpha16(img: &mut Image, y: RowsIter<[u8; 2]>, conv: RGBConvert<u16>, premultiplied_alpha: bool) -> Result<()> {
122    if let RGBConvert::Matrix(_) = conv {
123        return Err(Error::Unsupported("alpha image has color info"));
124    }
125    match img {
126        Image::RGBA8(img) => {
127            for (y_row, img_row) in y.zip(img.rows_mut()) {
128                if y_row.len() != img_row.len() {
129                    return Err(Error::Unsupported("invalid alpha size"));
130                }
131                for (y, px) in y_row.iter().copied().zip(img_row.iter_mut()) {
132                    let y = u16::from_ne_bytes(y);
133                    px.a = (conv.to_luma(y) >> 8) as u8;
134                }
135                if premultiplied_alpha {
136                    unpremultiply8(img_row);
137                }
138            }
139        },
140        Image::RGBA16(img) => {
141            for (y_row, img_row) in y.zip(img.rows_mut()) {
142                if y_row.len() != img_row.len() {
143                    return Err(Error::Unsupported("invalid alpha size"));
144                }
145                for (y, px) in y_row.iter().copied().zip(img_row.iter_mut()) {
146                    let y = u16::from_ne_bytes(y);
147                    px.a = conv.to_luma(y);
148                }
149                if premultiplied_alpha {
150                    unpremultiply16(img_row);
151                }
152            }
153        },
154        _ => return Err(Error::Unsupported("internal error")),
155    };
156    Ok(())
157}
158
159fn add_alpha8(img: &mut Image, y: RowsIter<u8>, conv: RGBConvert, premultiplied_alpha: bool) -> Result<()> {
160    if let RGBConvert::Matrix(_) = conv {
161        return Err(Error::Unsupported("alpha image has color info"));
162    }
163    match img {
164        Image::RGBA8(img) => {
165            for (y_row, img_row) in y.zip(img.rows_mut()) {
166                if y_row.len() != img_row.len() {
167                    return Err(Error::Unsupported("invalid alpha size"));
168                }
169                for (y, px) in y_row.iter().copied().zip(img_row.iter_mut()) {
170                    px.a = conv.to_luma(y);
171                }
172                if premultiplied_alpha {
173                    unpremultiply8(img_row);
174                }
175            }
176        },
177        Image::RGBA16(img) => {
178            for (y_row, img_row) in y.zip(img.rows_mut()) {
179                if y_row.len() != img_row.len() {
180                    return Err(Error::Unsupported("invalid alpha size"));
181                }
182                for (y, px) in y_row.iter().copied().zip(img_row.iter_mut()) {
183                    let g = u16::from(conv.to_luma(y));
184                    px.a = (g << 8) | g;
185                }
186                if premultiplied_alpha {
187                    unpremultiply16(img_row);
188                }
189            }
190        },
191        _ => return Err(Error::Unsupported("internal error")),
192    };
193    Ok(())
194}
195
196#[inline(never)]
197fn unpremultiply8(img_row: &mut [RGBA<u8>]) {
198    for px in img_row.iter_mut() {
199        if px.a != 255 && px.a != 0 {
200            *px.rgb_mut() = px.rgb().map(|c| (c as u16 * 255 / px.a as u16).min(255) as u8)
201        }
202    }
203}
204
205#[inline(never)]
206fn unpremultiply16(img_row: &mut [RGBA<u16>]) {
207    for px in img_row.iter_mut() {
208        if px.a != 0xFFFF && px.a != 0 {
209            *px.rgb_mut() = px.rgb().map(|c| (c as u32 * 0xFFFF / px.a as u32).min(0xFFFF) as u16)
210        }
211    }
212}
213
214fn yuv_to_rgb16(color: &FrameTempRef, range: Range, depth: Depth, y: RowsIter<[u8; 2]>, chroma_sampling: ChromaSampling, u: RowsIter<[u8; 2]>, v: RowsIter<[u8; 2]>, has_alpha: bool) -> Result<Image, Error> {
215    let mc = color.matrix_coefficients().unwrap_or(MatrixCoefficients::BT601);
216    let conv = RGBConvert::<u16>::new(range, mc, depth).inspect_err(move |e| {
217        log_unsupported_color(e, mc, range);
218    })?;
219    let width = y.width();
220    let height = y.height();
221    let mut tmp1;
222    let mut tmp2;
223    let mut tmp3;
224    let px_iter: &mut dyn Iterator<Item = YUV<[u8; 2]>> = match chroma_sampling {
225        ChromaSampling::Cs444 => {
226            tmp1 = yuv_444(y, u, v);
227            &mut tmp1
228        },
229        ChromaSampling::Cs420 => {
230            tmp2 = yuv_420(y, u, v);
231            &mut tmp2
232        },
233        ChromaSampling::Cs422 => {
234            tmp3 = yuv_422(y, u, v);
235            &mut tmp3
236        },
237        ChromaSampling::Monochrome => unreachable!(),
238    };
239    if has_alpha {
240        let mut out = Vec::with_capacity(width * height);
241        out.extend(px_iter.map(|px| conv.to_rgb(YUV{
242            y: u16::from_ne_bytes(px.y),
243            u: u16::from_ne_bytes(px.u),
244            v: u16::from_ne_bytes(px.v),
245        }).with_alpha(0)));
246        Ok(Image::RGBA16(ImgVec::new(out, width, height)))
247    } else {
248        let mut out = Vec::with_capacity(width * height);
249        out.extend(px_iter.map(|px| conv.to_rgb(YUV{
250            y: u16::from_ne_bytes(px.y),
251            u: u16::from_ne_bytes(px.u),
252            v: u16::from_ne_bytes(px.v),
253        })));
254        Ok(Image::RGB16(ImgVec::new(out, width, height)))
255    }
256}
257
258fn yuv_to_gray16(color: &FrameTempRef, range: Range, depth: Depth, y: RowsIter<[u8; 2]>, has_alpha: bool) -> Result<Image, Error> {
259    let mc = color.matrix_coefficients().unwrap_or(MatrixCoefficients::Identity);
260    let conv = RGBConvert::<u16>::new(range, mc, depth).inspect_err(move |e| {
261        log_unsupported_color(e, mc, range);
262    })?;
263    let width = y.width();
264    let height = y.height();
265    if has_alpha {
266        let mut out = Vec::with_capacity(width * height);
267        out.extend(y.flat_map(|row| {
268            row.iter().copied().map(|y| {
269                let g = conv.to_luma(u16::from_ne_bytes(y));
270                RGBA::new(g, g, g, 0)
271            })
272        }));
273        Ok(Image::RGBA16(ImgVec::new(out, width, height)))
274    } else {
275        let mut out = Vec::with_capacity(width * height);
276        out.extend(y.flat_map(|row| {
277            row.iter()
278                .copied()
279                .map(|y| conv.to_luma(u16::from_ne_bytes(y)))
280        }));
281        Ok(Image::Gray16(ImgVec::new(out, width, height)))
282    }
283}
284
285fn yuv_to_gray8(color: &FrameTempRef, range: Range, y: RowsIter<u8>, has_alpha: bool) -> Result<Image, Error> {
286    let mc = color.matrix_coefficients().unwrap_or(MatrixCoefficients::Identity);
287    let conv = RGBConvert::<u8>::new(range, mc).inspect_err(move |e| {
288        log_unsupported_color(e, mc, range);
289    })?;
290    let width = y.width();
291    let height = y.height();
292    if has_alpha {
293        let mut out = Vec::with_capacity(width * height);
294        out.extend(y.flat_map(|row| {
295            row.iter()
296                .copied()
297                .map(|y| conv.to_rgb(YUV { y, u: 128, v: 128 }).with_alpha(0))
298        }));
299        Ok(Image::RGBA8(ImgVec::new(out, width, height)))
300    } else {
301        let mut out = Vec::with_capacity(width * height);
302        out.extend(y.flat_map(|row| {
303            row.iter()
304                .copied()
305                .map(|y| conv.to_rgb(YUV { y, u: 128, v: 128 }).g)
306        }));
307        Ok(Image::Gray8(ImgVec::new(out, width, height)))
308    }
309}
310
311fn yuv_to_rgb8(color: &FrameTempRef, range: Range, y: RowsIter<u8>, chroma_sampling: ChromaSampling, u: RowsIter<u8>, v: RowsIter<u8>, has_alpha: bool) -> Result<Image, Error> {
312    let mc = color.matrix_coefficients().unwrap_or(MatrixCoefficients::BT601);
313    let conv = RGBConvert::<u8>::new(range, mc).inspect_err(move |e| {
314        log_unsupported_color(e, mc, range);
315    })?;
316    let width = y.width();
317    let height = y.height();
318    let mut tmp1;
319    let mut tmp2;
320    let mut tmp3;
321    let px_iter: &mut dyn Iterator<Item = YUV<u8>> = match chroma_sampling {
322        ChromaSampling::Cs444 => {
323            tmp1 = yuv_444(y, u, v);
324            &mut tmp1
325        },
326        ChromaSampling::Cs420 => {
327            tmp2 = yuv_420(y, u, v);
328            &mut tmp2
329        },
330        ChromaSampling::Cs422 => {
331            tmp3 = yuv_422(y, u, v);
332            &mut tmp3
333        },
334        ChromaSampling::Monochrome => unreachable!(),
335    };
336    if has_alpha {
337        let mut out = Vec::with_capacity(width * height);
338        out.extend(px_iter.map(|px| conv.to_rgb(px).with_alpha(0)));
339        Ok(Image::RGBA8(ImgVec::new(out, width, height)))
340    } else {
341        let mut out = Vec::with_capacity(width * height);
342        out.extend(px_iter.map(|px| conv.to_rgb(px)));
343        Ok(Image::RGB8(ImgVec::new(out, width, height)))
344    }
345}
346
347#[cold]
348fn log_unsupported_color(e: &yuv::Error, mc: MatrixCoefficients, range: Range) {
349    log::debug!("{e} (color={mc:?}, range={range:?})");
350}
aom_decode/avif.rs

aom_decode/
avif.rs
