1use aom_decode::chroma::{yuv_420, yuv_422, yuv_444};
2use aom_decode::color;
3use aom_decode::Config;
4use aom_decode::FrameTempRef;
5use aom_decode::RowsIters;
6use imgref::ImgVec;
7use rgb::prelude::*;
8use rgb::*;
9use std::io;
10use yuv::YUV;
11
12use quick_error::quick_error;
13
14quick_error! {
15 #[derive(Debug)]
16 pub enum Error {
17 Io(err: io::Error) {
18 display("{}", err)
19 from()
20 }
21 Parse(err: avif_parse::Error) {
22 display("{}", err)
23 from()
24 }
25 Decode(err: aom_decode::Error) {
26 display("{}", err)
27 from()
28 }
29 Meta(err: yuv::Error) {
30 display("{}", err)
31 from()
32 }
33 }
34}
35
36pub type Result<T, E = Error> = std::result::Result<T, E>;
37
38pub enum Image {
39 Rgb8(ImgVec<Rgb<u8>>),
40 Rgb16(ImgVec<Rgb<u16>>),
41 Rgba8(ImgVec<Rgba<u8>>),
42 Rgba16(ImgVec<Rgba<u16>>),
43 Gray8(ImgVec<Gray<u8>>),
44 Gray16(ImgVec<Gray<u16>>),
45}
46
47enum AlphaImage {
48 Gray8(ImgVec<Gray<u8>>),
49 Gray16(ImgVec<Gray<u16>>),
50}
51
52pub struct Decoder {
53 _decoder: Box<aom_decode::Decoder>,
54 color: FrameTempRef<'static>,
55 alpha: Option<AlphaImage>,
56 premultiplied_alpha: bool,
57}
58
59impl Decoder {
60 #[inline(always)]
61 pub fn from_avif(mut data: &[u8]) -> Result<Self> {
62 Self::from_reader(&mut data)
63 }
64
65 #[inline]
66 pub fn from_reader<R: io::Read>(reader: &mut R) -> Result<Self> {
67 let avif = avif_parse::read_avif(reader)?;
68 Self::from_parsed(avif)
69 }
70
71 fn from_parsed(avif: avif_parse::AvifData) -> Result<Self> {
72 let mut decoder = Box::new(aom_decode::Decoder::new(&Config {
73 threads: std::thread::available_parallelism().map_or(4, |a| a.get()).min(32),
74 })?);
75
76 let alpha = avif.alpha_item.as_ref().map(|a| Self::to_alpha(decoder.decode_frame(a)?)).transpose()?;
77 let premultiplied_alpha = avif.premultiplied_alpha;
78 let color = decoder.decode_frame(&avif.primary_item)?;
79 let color = unsafe {
82 std::mem::transmute::<FrameTempRef<'_>, FrameTempRef<'static>>(color)
83 };
84 Ok(Self {
85 _decoder: decoder,
86 color,
87 alpha,
88 premultiplied_alpha,
89 })
90 }
91
92 pub fn to_image(self) -> Result<Image> {
93 let color = Self::color_convert(self.color)?;
94 Ok(if let Some(alpha) = self.alpha {
95 let mut image = match (color, alpha) {
96 (Image::Rgb8(img), AlphaImage::Gray8(alpha)) => {
97 let buf = img.pixels().zip(alpha.pixels()).map(|(c, a)| c.with_alpha(*a)).collect();
98 Image::Rgba8(ImgVec::new(buf, img.width(), img.height()))
99 },
100 (Image::Rgb8(img), AlphaImage::Gray16(alpha)) => {
101 let buf = img.pixels().zip(alpha.pixels()).map(|(c, a)| c.map(|c| (u16::from(c) << 8) | u16::from(c)).with_alpha(*a)).collect();
102 Image::Rgba16(ImgVec::new(buf, img.width(), img.height()))
103 },
104 (Image::Rgb16(img), AlphaImage::Gray8(alpha)) => {
105 let buf = img.pixels().zip(alpha.pixels()).map(|(c, a)| c.with_alpha(u16::from(*a) << 8 | u16::from(*a))).collect();
106 Image::Rgba16(ImgVec::new(buf, img.width(), img.height()))
107 },
108 (Image::Rgb16(img), AlphaImage::Gray16(alpha)) => {
109 let buf = img.pixels().zip(alpha.pixels()).map(|(c, a)| c.with_alpha(*a)).collect();
110 Image::Rgba16(ImgVec::new(buf, img.width(), img.height()))
111 },
112 (Image::Rgba8(img), AlphaImage::Gray8(alpha)) => {
113 let buf = img.pixels().zip(alpha.pixels()).map(|(c, a)| c.with_alpha(*a)).collect();
114 Image::Rgba8(ImgVec::new(buf, img.width(), img.height()))
115 },
116 (Image::Gray8(img), AlphaImage::Gray8(alpha)) => {
117 let buf = img.pixels().zip(alpha.pixels()).map(|(c, a)| Rgba::new(*c,*c,*c,*a)).collect();
118 Image::Rgba8(ImgVec::new(buf, img.width(), img.height()))
119 },
120 (Image::Gray8(img), AlphaImage::Gray16(alpha)) => {
121 let buf = img.pixels().zip(alpha.pixels()).map(|(c, a)| {
122 let c = u16::from(*c) << 8 | u16::from(*c);
123 Rgba::new(c,c,c,*a)
124 }).collect();
125 Image::Rgba16(ImgVec::new(buf, img.width(), img.height()))
126 },
127 (Image::Gray16(img), AlphaImage::Gray8(alpha)) => {
128 let buf = img.pixels().zip(alpha.pixels()).map(|(c, a)| Rgba::new(*c,*c,*c,u16::from(*a) << 8 | u16::from(*a))).collect();
129 Image::Rgba16(ImgVec::new(buf, img.width(), img.height()))
130 },
131 (Image::Gray16(img), AlphaImage::Gray16(alpha)) => {
132 let buf = img.pixels().zip(alpha.pixels()).map(|(c, a)| Rgba::new(*c,*c,*c,*a)).collect();
133 Image::Rgba16(ImgVec::new(buf, img.width(), img.height()))
134 },
135 (Image::Rgba8(_) | Image::Rgba16(_), _) => unreachable!(),
136 };
137 if self.premultiplied_alpha {
138 match &mut image {
139 Image::Rgba8(img) => {
140 img.pixels_mut().filter(|px| px.a > 0).for_each(|px| {
141 #[inline(always)]
142 fn unprem(val: u8, alpha: u8) -> u8 {
143 ((u16::from(val) * 256) / (u16::from(alpha) * 256) / 256).min(255) as u8
144 }
145 px.r = unprem(px.r, px.a);
146 px.g = unprem(px.g, px.a);
147 px.b = unprem(px.b, px.a);
148 });
149 },
150 Image::Rgba16(img) => {
151 img.pixels_mut().filter(|px| px.a > 0).for_each(|px| {
152 #[inline(always)]
153 fn unprem(val: u16, alpha: u16) -> u16 {
154 ((u32::from(val) * 0xFFFF) / (u32::from(alpha) * 0xFFFF) / 0xFFFF).min(65535) as u16
155 }
156 px.r = unprem(px.r, px.a);
157 px.g = unprem(px.g, px.a);
158 px.b = unprem(px.b, px.a);
159 });
160 },
161 _ => {},
162 }
163 }
164 image
165 } else {
166 color
167 })
168 }
169
170 fn color_convert(img: FrameTempRef<'_>) -> Result<Image> {
171 let range = img.range();
172 Ok(match img.rows_iter()? {
173 RowsIters::YuvPlanes8 {y,u,v,chroma_sampling} => {
174 let mc = img.matrix_coefficients().unwrap_or(color::MatrixCoefficients::BT709);
175 let conv = yuv::convert::RGBConvert::<u8>::new(range, mc)?;
176 let width = y.width();
177 let height = y.height();
178 let mut out = Vec::with_capacity(width * height);
179 let mut tmp1;
180 let mut tmp2;
181 let mut tmp3;
182 let px_iter: &mut dyn Iterator<Item=YUV<u8>> = match chroma_sampling {
183 color::ChromaSampling::Cs444 => {
184 tmp1 = yuv_444(y, u, v);
185 &mut tmp1
186 },
187 color::ChromaSampling::Cs420 => {
188 tmp2 = yuv_420(y, u, v);
189 &mut tmp2
190 },
191 color::ChromaSampling::Cs422 => {
192 tmp3 = yuv_422(y, u, v);
193 &mut tmp3
194 },
195 color::ChromaSampling::Monochrome => return Err(Error::Meta(yuv::Error::InvalidDepthRequested)),
196 };
197 out.extend(px_iter.map(|px| conv.to_rgb(px)));
198 Image::Rgb8(ImgVec::new(out, width, height))
199 },
200 RowsIters::YuvPlanes16 {y,u,v,chroma_sampling, depth} => {
201 let mc = img.matrix_coefficients().unwrap_or(color::MatrixCoefficients::BT709);
202 let conv = yuv::convert::RGBConvert::<u16>::new(range, mc, depth)?;
203 let width = y.width();
204 let height = y.height();
205 let mut out = Vec::with_capacity(width * height);
206 let mut tmp1;
207 let mut tmp2;
208 let mut tmp3;
209 let px_iter: &mut dyn Iterator<Item=YUV<[u8; 2]>> = match chroma_sampling {
210 color::ChromaSampling::Cs444 => {
211 tmp1 = yuv_444(y, u, v);
212 &mut tmp1
213 },
214 color::ChromaSampling::Cs420 => {
215 tmp2 = yuv_420(y, u, v);
216 &mut tmp2
217 },
218 color::ChromaSampling::Cs422 => {
219 tmp3 = yuv_422(y, u, v);
220 &mut tmp3
221 },
222 color::ChromaSampling::Monochrome => unreachable!(),
223 };
224 out.extend(px_iter.map(|px| conv.to_rgb(YUV{
225 y: u16::from_ne_bytes(px.y),
226 u: u16::from_ne_bytes(px.u),
227 v: u16::from_ne_bytes(px.v),
228 })));
229 Image::Rgb16(ImgVec::new(out, width, height))
230 },
231 gray_iters => {
232 let mc = img.matrix_coefficients().unwrap_or(color::MatrixCoefficients::Identity);
233 match Self::to_gray(range, mc, gray_iters)? {
234 AlphaImage::Gray8(img) => Image::Gray8(img),
235 AlphaImage::Gray16(img) => Image::Gray16(img),
236 }
237 },
238 })
239 }
240
241 fn to_alpha(img: FrameTempRef<'_>) -> Result<AlphaImage> {
242 let range = img.range();
243 let mc = img.matrix_coefficients().unwrap_or(color::MatrixCoefficients::Identity);
244 Ok(Self::to_gray(range, mc, img.rows_iter()?)?)
245 }
246
247 fn to_gray(range: aom_decode::color::Range, mc: color::MatrixCoefficients, iters: RowsIters) -> Result<AlphaImage, yuv::Error> {
248 Ok(match iters {
249 RowsIters::YuvPlanes8 {y,..} | RowsIters::Mono8(y) => {
250 let conv = yuv::convert::RGBConvert::<u8>::new(range, mc)?;
251 let width = y.width();
252 let height = y.height();
253 let mut out = Vec::with_capacity(width * height);
254 out.extend(y.flat_map(|row| {
255 row.iter().copied().map(|y| {
256 Gray::new(conv.to_rgb(YUV{y,u:128,v:128}).g)
257 })
258 }));
259 AlphaImage::Gray8(ImgVec::new(out, width, height))
260 },
261 RowsIters::YuvPlanes16 {y, depth, ..} | RowsIters::Mono16(y, depth) => {
262 let conv = yuv::convert::RGBConvert::<u16>::new(range, mc, depth)?;
263 let width = y.width();
264 let height = y.height();
265 let mut out = Vec::with_capacity(width * height);
266 out.extend(y.flat_map(|row| {
267 row.iter().copied().map(|y| {
268 let y = u16::from_ne_bytes(y);
269 Gray::new(conv.to_rgb(YUV{y,u:128*256+128,v:128*256+128}).g)
270 })
271 }));
272 AlphaImage::Gray16(ImgVec::new(out, width, height))
273 },
274 })
275 }
276}