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