image/codecs/avif/
encoder.rs

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/
6use std::borrow::Cow;
7use std::cmp::min;
8use std::io::Write;
9use std::mem::size_of;
10
11use crate::buffer::ConvertBuffer;
12use crate::color::{FromColor, Luma, LumaA, Rgb, Rgba};
13use crate::error::{
14    EncodingError, ParameterError, ParameterErrorKind, UnsupportedError, UnsupportedErrorKind,
15};
16use crate::{ExtendedColorType, ImageBuffer, ImageEncoder, ImageFormat, Pixel};
17use crate::{ImageError, ImageResult};
18
19use bytemuck::{try_cast_slice, try_cast_slice_mut, Pod, PodCastError};
20use num_traits::Zero;
21use ravif::{Encoder, Img, RGB8, RGBA8};
22use rgb::AsPixels;
23
24/// AVIF Encoder.
25///
26/// Writes one image into the chosen output.
27pub 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]
35pub enum ColorSpace {
36    /// sRGB colorspace
37    Srgb,
38    /// BT.709 colorspace
39    Bt709,
40}
41
42impl ColorSpace {
43    fn to_ravif(self) -> ravif::ColorSpace {
44        match self {
45            Self::Srgb => ravif::ColorSpace::RGB,
46            Self::Bt709 => ravif::ColorSpace::YCbCr,
47        }
48    }
49}
50
51enum RgbColor<'buf> {
52    Rgb8(Img<&'buf [RGB8]>),
53    Rgba8(Img<&'buf [RGBA8]>),
54}
55
56impl<W: Write> AvifEncoder<W> {
57    /// Create a new encoder that writes its output to `w`.
58    pub fn new(w: W) -> Self {
59        AvifEncoder::new_with_speed_quality(w, 4, 80) // `cavif` uses these defaults
60    }
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    pub fn new_with_speed_quality(w: W, speed: u8, quality: u8) -> Self {
67        // Clamp quality and speed to range
68        let quality = min(quality, 100);
69        let speed = min(speed, 10);
70
71        let encoder = Encoder::new()
72            .with_quality(f32::from(quality))
73            .with_alpha_quality(f32::from(quality))
74            .with_speed(speed)
75            .with_depth(Some(8));
76
77        AvifEncoder { inner: w, encoder }
78    }
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_space(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
96impl<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    fn write_image(
104        mut self,
105        data: &[u8],
106        width: u32,
107        height: u32,
108        color: ExtendedColorType,
109    ) -> ImageResult<()> {
110        let expected_buffer_len = color.buffer_size(width, height);
111        assert_eq!(
112            expected_buffer_len,
113            data.len() as u64,
114            "Invalid buffer length: expected {expected_buffer_len} got {} for {width}x{height} image",
115            data.len(),
116        );
117
118        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        let mut fallback = vec![]; // This vector is used if we need to do a color conversion.
123        let result = match Self::encode_as_img(&mut fallback, data, width, height, color)? {
124            RgbColor::Rgb8(buffer) => self.encoder.encode_rgb(buffer),
125            RgbColor::Rgba8(buffer) => self.encoder.encode_rgba(buffer),
126        };
127        let data = result.map_err(|err| {
128            ImageError::Encoding(EncodingError::new(ImageFormat::Avif.into(), err))
129        })?;
130        self.inner.write_all(&data.avif_file)?;
131        Ok(())
132    }
133}
134
135impl<W: Write> AvifEncoder<W> {
136    // Does not currently do anything. Mirrors behaviour of old config function.
137    fn set_color(&mut self, _color: ExtendedColorType) {
138        // self.config.color_space = ColorSpace::RGB;
139    }
140
141    fn encode_as_img<'buf>(
142        fallback: &'buf mut Vec<u8>,
143        data: &'buf [u8],
144        width: u32,
145        height: u32,
146        color: ExtendedColorType,
147    ) -> ImageResult<RgbColor<'buf>> {
148        // Error wrapping utility for color dependent buffer dimensions.
149        fn try_from_raw<P: Pixel + 'static>(
150            data: &[P::Subpixel],
151            width: u32,
152            height: u32,
153        ) -> ImageResult<ImageBuffer<P, &[P::Subpixel]>> {
154            ImageBuffer::from_raw(width, height, data).ok_or_else(|| {
155                ImageError::Parameter(ParameterError::from_kind(
156                    ParameterErrorKind::DimensionMismatch,
157                ))
158            })
159        }
160
161        // Convert to target color type using few buffer allocations.
162        fn convert_into<'buf, P>(
163            buf: &'buf mut Vec<u8>,
164            image: ImageBuffer<P, &[P::Subpixel]>,
165        ) -> Img<&'buf [RGBA8]>
166        where
167            P: Pixel + 'static,
168            Rgba<u8>: FromColor<P>,
169        {
170            let (width, height) = image.dimensions();
171            // TODO: conversion re-using the target buffer?
172            let image: ImageBuffer<Rgba<u8>, _> = image.convert();
173            *buf = image.into_raw();
174            Img::new(buf.as_pixels(), width as usize, height as usize)
175        }
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        fn cast_buffer<Channel>(buf: &[u8]) -> ImageResult<Cow<[Channel]>>
180        where
181            Channel: Pod + Zero,
182        {
183            match try_cast_slice(buf) {
184                Ok(slice) => Ok(Cow::Borrowed(slice)),
185                Err(PodCastError::OutputSliceWouldHaveSlop) => Err(ImageError::Parameter(
186                    ParameterError::from_kind(ParameterErrorKind::DimensionMismatch),
187                )),
188                Err(PodCastError::TargetAlignmentGreaterAndInputNotAligned) => {
189                    // Sad, but let's allocate.
190                    // bytemuck checks alignment _before_ slop but size mismatch before this..
191                    if buf.len() % size_of::<Channel>() != 0 {
192                        Err(ImageError::Parameter(ParameterError::from_kind(
193                            ParameterErrorKind::DimensionMismatch,
194                        )))
195                    } else {
196                        let len = buf.len() / size_of::<Channel>();
197                        let mut data = vec![Channel::zero(); len];
198                        let view = try_cast_slice_mut::<_, u8>(data.as_mut_slice()).unwrap();
199                        view.copy_from_slice(buf);
200                        Ok(Cow::Owned(data))
201                    }
202                }
203                Err(err) => {
204                    // Are you trying to encode a ZST??
205                    Err(ImageError::Parameter(ParameterError::from_kind(
206                        ParameterErrorKind::Generic(format!("{err:?}")),
207                    )))
208                }
209            }
210        }
211
212        match color {
213            ExtendedColorType::Rgb8 => {
214                // ravif doesn't do any checks but has some asserts, so we do the checks.
215                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                if img.pixels().len() == 0 {
219                    return Err(ImageError::Parameter(ParameterError::from_kind(
220                        ParameterErrorKind::DimensionMismatch,
221                    )));
222                }
223
224                Ok(RgbColor::Rgb8(Img::new(
225                    AsPixels::as_pixels(data),
226                    width as usize,
227                    height as usize,
228                )))
229            }
230            ExtendedColorType::Rgba8 => {
231                // ravif doesn't do any checks but has some asserts, so we do the checks.
232                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                if img.pixels().len() == 0 {
236                    return Err(ImageError::Parameter(ParameterError::from_kind(
237                        ParameterErrorKind::DimensionMismatch,
238                    )));
239                }
240
241                Ok(RgbColor::Rgba8(Img::new(
242                    AsPixels::as_pixels(data),
243                    width as usize,
244                    height as usize,
245                )))
246            }
247            // we need a separate buffer..
248            ExtendedColorType::L8 => {
249                let image = try_from_raw::<Luma<u8>>(data, width, height)?;
250                Ok(RgbColor::Rgba8(convert_into(fallback, image)))
251            }
252            ExtendedColorType::La8 => {
253                let image = try_from_raw::<LumaA<u8>>(data, width, height)?;
254                Ok(RgbColor::Rgba8(convert_into(fallback, image)))
255            }
256            // we need to really convert data..
257            ExtendedColorType::L16 => {
258                let buffer = cast_buffer(data)?;
259                let image = try_from_raw::<Luma<u16>>(&buffer, width, height)?;
260                Ok(RgbColor::Rgba8(convert_into(fallback, image)))
261            }
262            ExtendedColorType::La16 => {
263                let buffer = cast_buffer(data)?;
264                let image = try_from_raw::<LumaA<u16>>(&buffer, width, height)?;
265                Ok(RgbColor::Rgba8(convert_into(fallback, image)))
266            }
267            ExtendedColorType::Rgb16 => {
268                let buffer = cast_buffer(data)?;
269                let image = try_from_raw::<Rgb<u16>>(&buffer, width, height)?;
270                Ok(RgbColor::Rgba8(convert_into(fallback, image)))
271            }
272            ExtendedColorType::Rgba16 => {
273                let buffer = cast_buffer(data)?;
274                let image = try_from_raw::<Rgba<u16>>(&buffer, width, height)?;
275                Ok(RgbColor::Rgba8(convert_into(fallback, image)))
276            }
277            // for cases we do not support at all?
278            _ => Err(ImageError::Unsupported(
279                UnsupportedError::from_format_and_kind(
280                    ImageFormat::Avif.into(),
281                    UnsupportedErrorKind::Color(color),
282                ),
283            )),
284        }
285    }
286}
image/codecs/avif/encoder.rs

image/codecs/avif/
encoder.rs
