Coverage Report

Created: 2025-11-09 06:56

next uncovered line (L), next uncovered region (R), next uncovered branch (B)
/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
}