/src/image/src/codecs/avif/encoder.rs
Line | Count | Source |
1 | | //! Encoding of AVIF images. |
2 | | /// |
3 | | /// The [AVIF] specification defines an image derivative of the AV1 bitstream, an open video codec. |
4 | | /// |
5 | | /// [AVIF]: https://aomediacodec.github.io/av1-avif/ |
6 | | use std::borrow::Cow; |
7 | | use std::cmp::min; |
8 | | use std::io::Write; |
9 | | use std::mem::size_of; |
10 | | |
11 | | use crate::color::{FromColor, Luma, LumaA, Rgb, Rgba}; |
12 | | use crate::error::{ |
13 | | EncodingError, ParameterError, ParameterErrorKind, UnsupportedError, UnsupportedErrorKind, |
14 | | }; |
15 | | use crate::{ExtendedColorType, ImageBuffer, ImageEncoder, ImageFormat, Pixel}; |
16 | | use crate::{ImageError, ImageResult}; |
17 | | |
18 | | use bytemuck::{try_cast_slice, try_cast_slice_mut, Pod, PodCastError}; |
19 | | use num_traits::Zero; |
20 | | use ravif::{BitDepth, Encoder, Img, RGB8, RGBA8}; |
21 | | use rgb::AsPixels; |
22 | | |
23 | | /// AVIF Encoder. |
24 | | /// |
25 | | /// Writes one image into the chosen output. |
26 | | pub struct AvifEncoder<W> { |
27 | | inner: W, |
28 | | encoder: Encoder<'static>, |
29 | | } |
30 | | |
31 | | /// An enumeration over supported AVIF color spaces |
32 | | #[derive(Debug, Copy, Clone, PartialEq, Eq)] |
33 | | #[non_exhaustive] |
34 | | pub enum ColorSpace { |
35 | | /// sRGB colorspace |
36 | | Srgb, |
37 | | /// BT.709 colorspace |
38 | | Bt709, |
39 | | } |
40 | | |
41 | | impl ColorSpace { |
42 | 0 | fn to_ravif(self) -> ravif::ColorModel { |
43 | 0 | match self { |
44 | 0 | Self::Srgb => ravif::ColorModel::RGB, |
45 | 0 | Self::Bt709 => ravif::ColorModel::YCbCr, |
46 | | } |
47 | 0 | } |
48 | | } |
49 | | |
50 | | enum RgbColor<'buf> { |
51 | | Rgb8(Img<&'buf [RGB8]>), |
52 | | Rgba8(Img<&'buf [RGBA8]>), |
53 | | } |
54 | | |
55 | | impl<W: Write> AvifEncoder<W> { |
56 | | /// Create a new encoder that writes its output to `w`. |
57 | 0 | pub fn new(w: W) -> Self { |
58 | 0 | AvifEncoder::new_with_speed_quality(w, 4, 80) // `cavif` uses these defaults |
59 | 0 | } Unexecuted instantiation: <image::codecs::avif::encoder::AvifEncoder<_>>::new Unexecuted instantiation: <image::codecs::avif::encoder::AvifEncoder<&mut std::io::cursor::Cursor<alloc::vec::Vec<u8>>>>::new |
60 | | |
61 | | /// Create a new encoder with a specified speed and quality that writes its output to `w`. |
62 | | /// `speed` accepts a value in the range 1-10, where 1 is the slowest and 10 is the fastest. |
63 | | /// Slower speeds generally yield better compression results. |
64 | | /// `quality` accepts a value in the range 1-100, where 1 is the worst and 100 is the best. |
65 | 0 | pub fn new_with_speed_quality(w: W, speed: u8, quality: u8) -> Self { |
66 | | // Clamp quality and speed to range |
67 | 0 | let quality = min(quality, 100); |
68 | 0 | let speed = min(speed, 10); |
69 | | |
70 | 0 | let encoder = Encoder::new() |
71 | 0 | .with_quality(f32::from(quality)) |
72 | 0 | .with_alpha_quality(f32::from(quality)) |
73 | 0 | .with_speed(speed) |
74 | 0 | .with_bit_depth(BitDepth::Eight); |
75 | | |
76 | 0 | AvifEncoder { inner: w, encoder } |
77 | 0 | } Unexecuted instantiation: <image::codecs::avif::encoder::AvifEncoder<_>>::new_with_speed_quality Unexecuted instantiation: <image::codecs::avif::encoder::AvifEncoder<&mut std::io::cursor::Cursor<alloc::vec::Vec<u8>>>>::new_with_speed_quality |
78 | | |
79 | | /// Encode with the specified `color_space`. |
80 | 0 | pub fn with_colorspace(mut self, color_space: ColorSpace) -> Self { |
81 | 0 | self.encoder = self |
82 | 0 | .encoder |
83 | 0 | .with_internal_color_model(color_space.to_ravif()); |
84 | 0 | self |
85 | 0 | } |
86 | | |
87 | | /// Configures `rayon` thread pool size. |
88 | | /// The default `None` is to use all threads in the default `rayon` thread pool. |
89 | 0 | pub fn with_num_threads(mut self, num_threads: Option<usize>) -> Self { |
90 | 0 | self.encoder = self.encoder.with_num_threads(num_threads); |
91 | 0 | self |
92 | 0 | } |
93 | | } |
94 | | |
95 | | impl<W: Write> ImageEncoder for AvifEncoder<W> { |
96 | | /// Encode image data with the indicated color type. |
97 | | /// |
98 | | /// The encoder currently requires all data to be RGBA8, it will be converted internally if |
99 | | /// necessary. When data is suitably aligned, i.e. u16 channels to two bytes, then the |
100 | | /// conversion may be more efficient. |
101 | | #[track_caller] |
102 | 0 | fn write_image( |
103 | 0 | mut self, |
104 | 0 | data: &[u8], |
105 | 0 | width: u32, |
106 | 0 | height: u32, |
107 | 0 | color: ExtendedColorType, |
108 | 0 | ) -> ImageResult<()> { |
109 | 0 | let expected_buffer_len = color.buffer_size(width, height); |
110 | 0 | assert_eq!( |
111 | | expected_buffer_len, |
112 | 0 | data.len() as u64, |
113 | 0 | "Invalid buffer length: expected {expected_buffer_len} got {} for {width}x{height} image", |
114 | 0 | data.len(), |
115 | | ); |
116 | | |
117 | 0 | self.set_color(color); |
118 | | // `ravif` needs strongly typed data so let's convert. We can either use a temporarily |
119 | | // owned version in our own buffer or zero-copy if possible by using the input buffer. |
120 | | // This requires going through `rgb`. |
121 | 0 | let mut fallback = vec![]; // This vector is used if we need to do a color conversion. |
122 | 0 | let result = match Self::encode_as_img(&mut fallback, data, width, height, color)? { |
123 | 0 | RgbColor::Rgb8(buffer) => self.encoder.encode_rgb(buffer), |
124 | 0 | RgbColor::Rgba8(buffer) => self.encoder.encode_rgba(buffer), |
125 | | }; |
126 | 0 | let data = result.map_err(|err| { |
127 | 0 | ImageError::Encoding(EncodingError::new(ImageFormat::Avif.into(), err)) |
128 | 0 | })?; Unexecuted instantiation: <image::codecs::avif::encoder::AvifEncoder<_> as image::io::encoder::ImageEncoder>::write_image::{closure#0}Unexecuted instantiation: <image::codecs::avif::encoder::AvifEncoder<&mut std::io::cursor::Cursor<alloc::vec::Vec<u8>>> as image::io::encoder::ImageEncoder>::write_image::{closure#0} |
129 | 0 | self.inner.write_all(&data.avif_file)?; |
130 | 0 | Ok(()) |
131 | 0 | } Unexecuted instantiation: <image::codecs::avif::encoder::AvifEncoder<_> as image::io::encoder::ImageEncoder>::write_image Unexecuted instantiation: <image::codecs::avif::encoder::AvifEncoder<&mut std::io::cursor::Cursor<alloc::vec::Vec<u8>>> as image::io::encoder::ImageEncoder>::write_image |
132 | | |
133 | 0 | fn set_exif_metadata(&mut self, exif: Vec<u8>) -> Result<(), UnsupportedError> { |
134 | | // encoder.with_exif() accepts Self rather than &mut self, and Encoder doesn't impl Default, |
135 | | // so we can't even mem::take it and have to do this instead |
136 | 0 | let encoder = std::mem::replace(&mut self.encoder, Encoder::new()); |
137 | 0 | self.encoder = encoder.with_exif(exif); |
138 | 0 | Ok(()) |
139 | 0 | } Unexecuted instantiation: <image::codecs::avif::encoder::AvifEncoder<_> as image::io::encoder::ImageEncoder>::set_exif_metadata Unexecuted instantiation: <image::codecs::avif::encoder::AvifEncoder<&mut std::io::cursor::Cursor<alloc::vec::Vec<u8>>> as image::io::encoder::ImageEncoder>::set_exif_metadata |
140 | | } |
141 | | |
142 | | impl<W: Write> AvifEncoder<W> { |
143 | | // Does not currently do anything. Mirrors behaviour of old config function. |
144 | 0 | fn set_color(&mut self, _color: ExtendedColorType) { |
145 | | // self.config.color_space = ColorSpace::RGB; |
146 | 0 | } Unexecuted instantiation: <image::codecs::avif::encoder::AvifEncoder<_>>::set_color Unexecuted instantiation: <image::codecs::avif::encoder::AvifEncoder<&mut std::io::cursor::Cursor<alloc::vec::Vec<u8>>>>::set_color |
147 | | |
148 | 0 | fn encode_as_img<'buf>( |
149 | 0 | fallback: &'buf mut Vec<u8>, |
150 | 0 | data: &'buf [u8], |
151 | 0 | width: u32, |
152 | 0 | height: u32, |
153 | 0 | color: ExtendedColorType, |
154 | 0 | ) -> ImageResult<RgbColor<'buf>> { |
155 | | // Error wrapping utility for color dependent buffer dimensions. |
156 | 0 | fn try_from_raw<P: Pixel>( |
157 | 0 | data: &[P::Subpixel], |
158 | 0 | width: u32, |
159 | 0 | height: u32, |
160 | 0 | ) -> ImageResult<ImageBuffer<P, &[P::Subpixel]>> { |
161 | 0 | ImageBuffer::from_raw(width, height, data).ok_or_else(|| { |
162 | 0 | ImageError::Parameter(ParameterError::from_kind( |
163 | 0 | ParameterErrorKind::DimensionMismatch, |
164 | 0 | )) |
165 | 0 | }) Unexecuted instantiation: <image::codecs::avif::encoder::AvifEncoder<_>>::encode_as_img::try_from_raw::<_>::{closure#0}Unexecuted instantiation: <image::codecs::avif::encoder::AvifEncoder<_>>::encode_as_img::try_from_raw::<image::color::Rgb<u8>>::{closure#0}Unexecuted instantiation: <image::codecs::avif::encoder::AvifEncoder<_>>::encode_as_img::try_from_raw::<image::color::Rgb<u16>>::{closure#0}Unexecuted instantiation: <image::codecs::avif::encoder::AvifEncoder<_>>::encode_as_img::try_from_raw::<image::color::Luma<u8>>::{closure#0}Unexecuted instantiation: <image::codecs::avif::encoder::AvifEncoder<_>>::encode_as_img::try_from_raw::<image::color::Luma<u16>>::{closure#0}Unexecuted instantiation: <image::codecs::avif::encoder::AvifEncoder<_>>::encode_as_img::try_from_raw::<image::color::Rgba<u8>>::{closure#0}Unexecuted instantiation: <image::codecs::avif::encoder::AvifEncoder<_>>::encode_as_img::try_from_raw::<image::color::Rgba<u16>>::{closure#0}Unexecuted instantiation: <image::codecs::avif::encoder::AvifEncoder<_>>::encode_as_img::try_from_raw::<image::color::LumaA<u8>>::{closure#0}Unexecuted instantiation: <image::codecs::avif::encoder::AvifEncoder<_>>::encode_as_img::try_from_raw::<image::color::LumaA<u16>>::{closure#0} |
166 | 0 | } Unexecuted instantiation: <image::codecs::avif::encoder::AvifEncoder<_>>::encode_as_img::try_from_raw::<_> Unexecuted instantiation: <image::codecs::avif::encoder::AvifEncoder<_>>::encode_as_img::try_from_raw::<image::color::Rgb<u8>> Unexecuted instantiation: <image::codecs::avif::encoder::AvifEncoder<_>>::encode_as_img::try_from_raw::<image::color::Rgb<u16>> Unexecuted instantiation: <image::codecs::avif::encoder::AvifEncoder<_>>::encode_as_img::try_from_raw::<image::color::Luma<u8>> Unexecuted instantiation: <image::codecs::avif::encoder::AvifEncoder<_>>::encode_as_img::try_from_raw::<image::color::Luma<u16>> Unexecuted instantiation: <image::codecs::avif::encoder::AvifEncoder<_>>::encode_as_img::try_from_raw::<image::color::Rgba<u8>> Unexecuted instantiation: <image::codecs::avif::encoder::AvifEncoder<_>>::encode_as_img::try_from_raw::<image::color::Rgba<u16>> Unexecuted instantiation: <image::codecs::avif::encoder::AvifEncoder<_>>::encode_as_img::try_from_raw::<image::color::LumaA<u8>> Unexecuted instantiation: <image::codecs::avif::encoder::AvifEncoder<_>>::encode_as_img::try_from_raw::<image::color::LumaA<u16>> |
167 | | |
168 | | // Convert to target color type using few buffer allocations. |
169 | 0 | fn convert_into<'buf, P>( |
170 | 0 | buf: &'buf mut Vec<u8>, |
171 | 0 | image: ImageBuffer<P, &[P::Subpixel]>, |
172 | 0 | ) -> Img<&'buf [RGBA8]> |
173 | 0 | where |
174 | 0 | P: Pixel, |
175 | 0 | Rgba<u8>: FromColor<P>, |
176 | | { |
177 | 0 | let (width, height) = image.dimensions(); |
178 | | // TODO: conversion re-using the target buffer? |
179 | 0 | let image: ImageBuffer<Rgba<u8>, _> = image.convert(); |
180 | 0 | *buf = image.into_raw(); |
181 | 0 | Img::new(buf.as_pixels(), width as usize, height as usize) |
182 | 0 | } Unexecuted instantiation: <image::codecs::avif::encoder::AvifEncoder<_>>::encode_as_img::convert_into::<_> Unexecuted instantiation: <image::codecs::avif::encoder::AvifEncoder<_>>::encode_as_img::convert_into::<image::color::Rgb<u16>> Unexecuted instantiation: <image::codecs::avif::encoder::AvifEncoder<_>>::encode_as_img::convert_into::<image::color::Luma<u8>> Unexecuted instantiation: <image::codecs::avif::encoder::AvifEncoder<_>>::encode_as_img::convert_into::<image::color::Luma<u16>> Unexecuted instantiation: <image::codecs::avif::encoder::AvifEncoder<_>>::encode_as_img::convert_into::<image::color::Rgba<u16>> Unexecuted instantiation: <image::codecs::avif::encoder::AvifEncoder<_>>::encode_as_img::convert_into::<image::color::LumaA<u8>> Unexecuted instantiation: <image::codecs::avif::encoder::AvifEncoder<_>>::encode_as_img::convert_into::<image::color::LumaA<u16>> |
183 | | |
184 | | // Cast the input slice using few buffer allocations if possible. |
185 | | // In particular try not to allocate if the caller did the infallible reverse. |
186 | 0 | fn cast_buffer<Channel>(buf: &[u8]) -> ImageResult<Cow<'_, [Channel]>> |
187 | 0 | where |
188 | 0 | Channel: Pod + Zero, |
189 | | { |
190 | 0 | match try_cast_slice(buf) { |
191 | 0 | Ok(slice) => Ok(Cow::Borrowed(slice)), |
192 | 0 | Err(PodCastError::OutputSliceWouldHaveSlop) => Err(ImageError::Parameter( |
193 | 0 | ParameterError::from_kind(ParameterErrorKind::DimensionMismatch), |
194 | 0 | )), |
195 | | Err(PodCastError::TargetAlignmentGreaterAndInputNotAligned) => { |
196 | | // Sad, but let's allocate. |
197 | | // bytemuck checks alignment _before_ slop but size mismatch before this.. |
198 | 0 | if !buf.len().is_multiple_of(size_of::<Channel>()) { |
199 | 0 | Err(ImageError::Parameter(ParameterError::from_kind( |
200 | 0 | ParameterErrorKind::DimensionMismatch, |
201 | 0 | ))) |
202 | | } else { |
203 | 0 | let len = buf.len() / size_of::<Channel>(); |
204 | 0 | let mut data = vec![Channel::zero(); len]; |
205 | 0 | let view = try_cast_slice_mut::<_, u8>(data.as_mut_slice()).unwrap(); |
206 | 0 | view.copy_from_slice(buf); |
207 | 0 | Ok(Cow::Owned(data)) |
208 | | } |
209 | | } |
210 | 0 | Err(err) => { |
211 | | // Are you trying to encode a ZST?? |
212 | 0 | Err(ImageError::Parameter(ParameterError::from_kind( |
213 | 0 | ParameterErrorKind::Generic(format!("{err:?}")), |
214 | 0 | ))) |
215 | | } |
216 | | } |
217 | 0 | } Unexecuted instantiation: <image::codecs::avif::encoder::AvifEncoder<_>>::encode_as_img::cast_buffer::<_> Unexecuted instantiation: <image::codecs::avif::encoder::AvifEncoder<_>>::encode_as_img::cast_buffer::<u16> |
218 | | |
219 | 0 | match color { |
220 | | ExtendedColorType::Rgb8 => { |
221 | | // ravif doesn't do any checks but has some asserts, so we do the checks. |
222 | 0 | let img = try_from_raw::<Rgb<u8>>(data, width, height)?; |
223 | | // Now, internally ravif uses u32 but it takes usize. We could do some checked |
224 | | // conversion but instead we use that a non-empty image must be addressable. |
225 | 0 | if img.pixels().is_empty() { |
226 | 0 | return Err(ImageError::Parameter(ParameterError::from_kind( |
227 | 0 | ParameterErrorKind::DimensionMismatch, |
228 | 0 | ))); |
229 | 0 | } |
230 | | |
231 | 0 | Ok(RgbColor::Rgb8(Img::new( |
232 | 0 | AsPixels::as_pixels(data), |
233 | 0 | width as usize, |
234 | 0 | height as usize, |
235 | 0 | ))) |
236 | | } |
237 | | ExtendedColorType::Rgba8 => { |
238 | | // ravif doesn't do any checks but has some asserts, so we do the checks. |
239 | 0 | let img = try_from_raw::<Rgba<u8>>(data, width, height)?; |
240 | | // Now, internally ravif uses u32 but it takes usize. We could do some checked |
241 | | // conversion but instead we use that a non-empty image must be addressable. |
242 | 0 | if img.pixels().is_empty() { |
243 | 0 | return Err(ImageError::Parameter(ParameterError::from_kind( |
244 | 0 | ParameterErrorKind::DimensionMismatch, |
245 | 0 | ))); |
246 | 0 | } |
247 | | |
248 | 0 | Ok(RgbColor::Rgba8(Img::new( |
249 | 0 | AsPixels::as_pixels(data), |
250 | 0 | width as usize, |
251 | 0 | height as usize, |
252 | 0 | ))) |
253 | | } |
254 | | // we need a separate buffer.. |
255 | | ExtendedColorType::L8 => { |
256 | 0 | let image = try_from_raw::<Luma<u8>>(data, width, height)?; |
257 | 0 | Ok(RgbColor::Rgba8(convert_into(fallback, image))) |
258 | | } |
259 | | ExtendedColorType::La8 => { |
260 | 0 | let image = try_from_raw::<LumaA<u8>>(data, width, height)?; |
261 | 0 | Ok(RgbColor::Rgba8(convert_into(fallback, image))) |
262 | | } |
263 | | // we need to really convert data.. |
264 | | ExtendedColorType::L16 => { |
265 | 0 | let buffer = cast_buffer(data)?; |
266 | 0 | let image = try_from_raw::<Luma<u16>>(&buffer, width, height)?; |
267 | 0 | Ok(RgbColor::Rgba8(convert_into(fallback, image))) |
268 | | } |
269 | | ExtendedColorType::La16 => { |
270 | 0 | let buffer = cast_buffer(data)?; |
271 | 0 | let image = try_from_raw::<LumaA<u16>>(&buffer, width, height)?; |
272 | 0 | Ok(RgbColor::Rgba8(convert_into(fallback, image))) |
273 | | } |
274 | | ExtendedColorType::Rgb16 => { |
275 | 0 | let buffer = cast_buffer(data)?; |
276 | 0 | let image = try_from_raw::<Rgb<u16>>(&buffer, width, height)?; |
277 | 0 | Ok(RgbColor::Rgba8(convert_into(fallback, image))) |
278 | | } |
279 | | ExtendedColorType::Rgba16 => { |
280 | 0 | let buffer = cast_buffer(data)?; |
281 | 0 | let image = try_from_raw::<Rgba<u16>>(&buffer, width, height)?; |
282 | 0 | Ok(RgbColor::Rgba8(convert_into(fallback, image))) |
283 | | } |
284 | | // for cases we do not support at all? |
285 | 0 | _ => Err(ImageError::Unsupported( |
286 | 0 | UnsupportedError::from_format_and_kind( |
287 | 0 | ImageFormat::Avif.into(), |
288 | 0 | UnsupportedErrorKind::Color(color), |
289 | 0 | ), |
290 | 0 | )), |
291 | | } |
292 | 0 | } Unexecuted instantiation: <image::codecs::avif::encoder::AvifEncoder<_>>::encode_as_img Unexecuted instantiation: <image::codecs::avif::encoder::AvifEncoder<&mut std::io::cursor::Cursor<alloc::vec::Vec<u8>>>>::encode_as_img |
293 | | } |