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