ravif/
av1encoder.rs

1#![allow(deprecated)]
2use std::borrow::Cow;
3use crate::dirtyalpha::blurred_dirty_alpha;
4use crate::error::Error;
5#[cfg(not(feature = "threading"))]
6use crate::rayoff as rayon;
7use imgref::{Img, ImgVec};
8use rav1e::prelude::*;
9use rgb::{RGB8, RGBA8};
10
11/// For [`Encoder::with_internal_color_model`]
12#[derive(Debug, Copy, Clone, Eq, PartialEq)]
13pub enum ColorModel {
14    /// Standard color model for photographic content. Usually the best choice.
15    /// This library always uses full-resolution color (4:4:4).
16    /// This library will automatically choose between BT.601 or BT.709.
17    YCbCr,
18    /// RGB channels are encoded without color space transformation.
19    /// Usually results in larger file sizes, and is less compatible than `YCbCr`.
20    /// Use only if the content really makes use of RGB, e.g. anaglyph images or RGB subpixel anti-aliasing.
21    RGB,
22}
23
24/// Handling of color channels in transparent images. For [`Encoder::with_alpha_color_mode`]
25#[derive(Debug, Copy, Clone, Eq, PartialEq)]
26pub enum AlphaColorMode {
27    /// Use unassociated alpha channel and leave color channels unchanged, even if there's redundant color data in transparent areas.
28    UnassociatedDirty,
29    /// Use unassociated alpha channel, but set color channels of transparent areas to a solid color to eliminate invisible data and improve compression.
30    UnassociatedClean,
31    /// Store color channels of transparent images in premultiplied form.
32    /// This requires support for premultiplied alpha in AVIF decoders.
33    ///
34    /// It may reduce file sizes due to clearing of fully-transparent pixels, but
35    /// may also increase file sizes due to creation of new edges in the color channels.
36    ///
37    /// Note that this is only internal detail for the AVIF file.
38    /// It does not change meaning of `RGBA` in this library — it's always unassociated.
39    Premultiplied,
40}
41
42#[derive(Default, Debug, Copy, Clone, Eq, PartialEq)]
43pub enum BitDepth {
44    Eight,
45    Ten,
46    /// Same as `Ten`
47    #[default]
48    Auto,
49}
50
51/// The newly-created image file + extra info FYI
52#[non_exhaustive]
53#[derive(Clone)]
54pub struct EncodedImage {
55    /// AVIF (HEIF+AV1) encoded image data
56    pub avif_file: Vec<u8>,
57    /// FYI: number of bytes of AV1 payload used for the color
58    pub color_byte_size: usize,
59    /// FYI: number of bytes of AV1 payload used for the alpha channel
60    pub alpha_byte_size: usize,
61}
62
63/// Encoder config builder
64///
65/// The lifetime is relevant only for [`Encoder::with_exif()`]. Use `Encoder<'static>` if Rust complains.
66#[derive(Debug, Clone)]
67pub struct Encoder<'exif_slice> {
68    /// 0-255 scale
69    quantizer: u8,
70    /// 0-255 scale
71    alpha_quantizer: u8,
72    /// rav1e preset 1 (slow) 10 (fast but crappy)
73    speed: u8,
74    /// True if RGBA input has already been premultiplied. It inserts appropriate metadata.
75    premultiplied_alpha: bool,
76    /// Which pixel format to use in AVIF file. RGB tends to give larger files.
77    color_model: ColorModel,
78    /// How many threads should be used (0 = match core count), None - use global rayon thread pool
79    threads: Option<usize>,
80    /// [`AlphaColorMode`]
81    alpha_color_mode: AlphaColorMode,
82    /// 8 or 10
83    output_depth: BitDepth,
84    /// Dropped into MPEG infe BOX
85    exif: Option<Cow<'exif_slice, [u8]>>,
86}
87
88/// Builder methods
89impl<'exif_slice> Encoder<'exif_slice> {
90    /// Start here
91    #[must_use]
92    pub fn new() -> Self {
93        Self {
94            quantizer: quality_to_quantizer(80.),
95            alpha_quantizer: quality_to_quantizer(80.),
96            speed: 5,
97            output_depth: BitDepth::default(),
98            premultiplied_alpha: false,
99            color_model: ColorModel::YCbCr,
100            threads: None,
101            exif: None,
102            alpha_color_mode: AlphaColorMode::UnassociatedClean,
103        }
104    }
105
106    /// Quality `1..=100`. Panics if out of range.
107    #[inline(always)]
108    #[track_caller]
109    #[must_use]
110    pub fn with_quality(mut self, quality: f32) -> Self {
111        assert!(quality >= 1. && quality <= 100.);
112        self.quantizer = quality_to_quantizer(quality);
113        self
114    }
115
116    #[doc(hidden)]
117    #[deprecated(note = "Renamed to with_bit_depth")]
118    #[must_use]
119    pub fn with_depth(self, depth: Option<u8>) -> Self {
120        self.with_bit_depth(depth.map(|d| if d >= 10 { BitDepth::Ten } else { BitDepth::Eight }).unwrap_or(BitDepth::Auto))
121    }
122
123    /// Internal precision to use in the encoded AV1 data, for both color and alpha. 10-bit depth works best, even for 8-bit inputs/outputs.
124    ///
125    /// Use 8-bit depth only as a workaround for decoders that need it.
126    ///
127    /// This setting does not affect pixel inputs for this library.
128    #[inline(always)]
129    #[must_use]
130    pub fn with_bit_depth(mut self, depth: BitDepth) -> Self {
131        self.output_depth = depth;
132        self
133    }
134
135    /// Quality for the alpha channel only. `1..=100`. Panics if out of range.
136    #[inline(always)]
137    #[track_caller]
138    #[must_use]
139    pub fn with_alpha_quality(mut self, quality: f32) -> Self {
140        assert!(quality >= 1. && quality <= 100.);
141        self.alpha_quantizer = quality_to_quantizer(quality);
142        self
143    }
144
145    /// * 1 = very very slow, but max compression.
146    /// * 10 = quick, but larger file sizes and lower quality.
147    ///
148    /// Panics if outside `1..=10`.
149    #[inline(always)]
150    #[track_caller]
151    #[must_use]
152    pub fn with_speed(mut self, speed: u8) -> Self {
153        assert!(speed >= 1 && speed <= 10);
154        self.speed = speed;
155        self
156    }
157
158    /// Changes how color channels are stored in the image. The default is YCbCr.
159    ///
160    /// Note that this is only internal detail for the AVIF file, and doesn't
161    /// change color model of inputs to encode functions.
162    #[inline(always)]
163    #[must_use]
164    pub fn with_internal_color_model(mut self, color_model: ColorModel) -> Self {
165        self.color_model = color_model;
166        self
167    }
168
169    #[doc(hidden)]
170    #[deprecated = "Renamed to `with_internal_color_model()`"]
171    #[must_use]
172    pub fn with_internal_color_space(self, color_model: ColorModel) -> Self {
173        self.with_internal_color_model(color_model)
174    }
175
176    /// Configures `rayon` thread pool size.
177    /// The default `None` is to use all threads in the default `rayon` thread pool.
178    #[inline(always)]
179    #[track_caller]
180    #[must_use]
181    pub fn with_num_threads(mut self, num_threads: Option<usize>) -> Self {
182        assert!(num_threads.is_none_or(|n| n > 0));
183        self.threads = num_threads;
184        self
185    }
186
187    /// Configure handling of color channels in transparent images
188    ///
189    /// Note that this doesn't affect input format for this library,
190    /// which for RGBA is always uncorrelated alpha.
191    #[inline(always)]
192    #[must_use]
193    pub fn with_alpha_color_mode(mut self, mode: AlphaColorMode) -> Self {
194        self.alpha_color_mode = mode;
195        self.premultiplied_alpha = mode == AlphaColorMode::Premultiplied;
196        self
197    }
198
199    /// Embedded into AVIF file as-is
200    ///
201    /// The data can be `Vec<u8>`, or `&[u8]` if the encoder instance doesn't leave its scope.
202    pub fn with_exif(mut self, exif_data: impl Into<Cow<'exif_slice, [u8]>>) -> Self {
203        self.exif = Some(exif_data.into());
204        self
205    }
206}
207
208/// Once done with config, call one of the `encode_*` functions
209impl Encoder<'_> {
210    /// Make a new AVIF image from RGBA pixels (non-premultiplied, alpha last)
211    ///
212    /// Make the `Img` for the `buffer` like this:
213    ///
214    /// ```rust,ignore
215    /// Img::new(&pixels_rgba[..], width, height)
216    /// ```
217    ///
218    /// If you have pixels as `u8` slice, then use the `rgb` crate, and do:
219    ///
220    /// ```rust,ignore
221    /// use rgb::ComponentSlice;
222    /// let pixels_rgba = pixels_u8.as_rgba();
223    /// ```
224    ///
225    /// If all pixels are opaque, the alpha channel will be left out automatically.
226    ///
227    /// This function takes 8-bit inputs, but will generate an AVIF file using 10-bit depth.
228    ///
229    /// returns AVIF file with info about sizes about AV1 payload.
230    pub fn encode_rgba(&self, in_buffer: Img<&[rgb::RGBA<u8>]>) -> Result<EncodedImage, Error> {
231        let new_alpha = self.convert_alpha_8bit(in_buffer);
232        let buffer = new_alpha.as_ref().map(|b| b.as_ref()).unwrap_or(in_buffer);
233        let use_alpha = buffer.pixels().any(|px| px.a != 255);
234        if !use_alpha {
235            return self.encode_rgb_internal_from_8bit(buffer.width(), buffer.height(), buffer.pixels().map(|px| px.rgb()));
236        }
237
238        let width = buffer.width();
239        let height = buffer.height();
240        let matrix_coefficients = match self.color_model {
241            ColorModel::YCbCr => MatrixCoefficients::BT601,
242            ColorModel::RGB => MatrixCoefficients::Identity,
243        };
244        match self.output_depth {
245            BitDepth::Eight => {
246                let planes = buffer.pixels().map(|px| match self.color_model {
247                    ColorModel::YCbCr => rgb_to_8_bit_ycbcr(px.rgb(), BT601).into(),
248                    ColorModel::RGB => rgb_to_8_bit_gbr(px.rgb()).into(),
249                });
250                let alpha = buffer.pixels().map(|px| px.a);
251                self.encode_raw_planes_8_bit(width, height, planes, Some(alpha), PixelRange::Full, matrix_coefficients)
252            },
253            BitDepth::Ten | BitDepth::Auto => {
254                let planes = buffer.pixels().map(|px| match self.color_model {
255                    ColorModel::YCbCr => rgb_to_10_bit_ycbcr(px.rgb(), BT601).into(),
256                    ColorModel::RGB => rgb_to_10_bit_gbr(px.rgb()).into(),
257                });
258                let alpha = buffer.pixels().map(|px| to_ten(px.a));
259                self.encode_raw_planes_10_bit(width, height, planes, Some(alpha), PixelRange::Full, matrix_coefficients)
260            },
261        }
262    }
263
264    fn convert_alpha_8bit(&self, in_buffer: Img<&[RGBA8]>) -> Option<ImgVec<RGBA8>> {
265        match self.alpha_color_mode {
266            AlphaColorMode::UnassociatedDirty => None,
267            AlphaColorMode::UnassociatedClean => blurred_dirty_alpha(in_buffer),
268            AlphaColorMode::Premultiplied => {
269                let prem = in_buffer.pixels()
270                    .map(|px| {
271                        if px.a == 0 || px.a == 255 {
272                            RGBA8::default()
273                        } else {
274                            RGBA8::new(
275                                (u16::from(px.r) * 255 / u16::from(px.a)) as u8,
276                                (u16::from(px.g) * 255 / u16::from(px.a)) as u8,
277                                (u16::from(px.b) * 255 / u16::from(px.a)) as u8,
278                                px.a,
279                            )
280                        }
281                    })
282                    .collect();
283                Some(ImgVec::new(prem, in_buffer.width(), in_buffer.height()))
284            },
285        }
286    }
287
288    /// Make a new AVIF image from RGB pixels
289    ///
290    /// Make the `Img` for the `buffer` like this:
291    ///
292    /// ```rust,ignore
293    /// Img::new(&pixels_rgb[..], width, height)
294    /// ```
295    ///
296    /// If you have pixels as `u8` slice, then first do:
297    ///
298    /// ```rust,ignore
299    /// use rgb::ComponentSlice;
300    /// let pixels_rgb = pixels_u8.as_rgb();
301    /// ```
302    ///
303    /// returns AVIF file, size of color metadata
304    #[inline]
305    pub fn encode_rgb(&self, buffer: Img<&[RGB8]>) -> Result<EncodedImage, Error> {
306        self.encode_rgb_internal_from_8bit(buffer.width(), buffer.height(), buffer.pixels())
307    }
308
309    fn encode_rgb_internal_from_8bit(&self, width: usize, height: usize, pixels: impl Iterator<Item = RGB8> + Send + Sync) -> Result<EncodedImage, Error> {
310        let matrix_coefficients = match self.color_model {
311            ColorModel::YCbCr => MatrixCoefficients::BT601,
312            ColorModel::RGB => MatrixCoefficients::Identity,
313        };
314
315        match self.output_depth {
316            BitDepth::Eight => {
317                let planes = pixels.map(|px| {
318                    let (y, u, v) = match self.color_model {
319                        ColorModel::YCbCr => rgb_to_8_bit_ycbcr(px, BT601),
320                        ColorModel::RGB => rgb_to_8_bit_gbr(px),
321                    };
322                    [y, u, v]
323                });
324                self.encode_raw_planes_8_bit(width, height, planes, None::<[_; 0]>, PixelRange::Full, matrix_coefficients)
325            },
326            BitDepth::Ten | BitDepth::Auto => {
327                let planes = pixels.map(|px| {
328                    let (y, u, v) = match self.color_model {
329                        ColorModel::YCbCr => rgb_to_10_bit_ycbcr(px, BT601),
330                        ColorModel::RGB => rgb_to_10_bit_gbr(px),
331                    };
332                    [y, u, v]
333                });
334                self.encode_raw_planes_10_bit(width, height, planes, None::<[_; 0]>, PixelRange::Full, matrix_coefficients)
335            },
336        }
337    }
338
339    /// Encodes AVIF from 3 planar channels that are in the color space described by `matrix_coefficients`,
340    /// with sRGB transfer characteristics and color primaries.
341    ///
342    /// Alpha always uses full range. Chroma subsampling is not supported, and it's a bad idea for AVIF anyway.
343    /// If there's no alpha, use `None::<[_; 0]>`.
344    ///
345    /// `color_pixel_range` should be `PixelRange::Full` to avoid worsening already small 8-bit dynamic range.
346    /// Support for limited range may be removed in the future.
347    ///
348    /// If `AlphaColorMode::Premultiplied` has been set, the alpha pixels must be premultiplied.
349    /// `AlphaColorMode::UnassociatedClean` has no effect in this function, and is equivalent to `AlphaColorMode::UnassociatedDirty`.
350    ///
351    /// returns AVIF file, size of color metadata, size of alpha metadata overhead
352    #[inline]
353    pub fn encode_raw_planes_8_bit(
354        &self, width: usize, height: usize,
355        planes: impl IntoIterator<Item = [u8; 3]> + Send,
356        alpha: Option<impl IntoIterator<Item = u8> + Send>,
357        color_pixel_range: PixelRange, matrix_coefficients: MatrixCoefficients,
358    ) -> Result<EncodedImage, Error> {
359        self.encode_raw_planes_internal(width, height, planes, alpha, color_pixel_range, matrix_coefficients, 8)
360    }
361
362    /// Encodes AVIF from 3 planar channels that are in the color space described by `matrix_coefficients`,
363    /// with sRGB transfer characteristics and color primaries.
364    ///
365    /// The pixels are 10-bit (values `0.=1023`) in host's native endian.
366    ///
367    /// Alpha always uses full range. Chroma subsampling is not supported, and it's a bad idea for AVIF anyway.
368    /// If there's no alpha, use `None::<[_; 0]>`.
369    ///
370    /// `color_pixel_range` should be `PixelRange::Full`. Support for limited range may be removed in the future.
371    ///
372    /// If `AlphaColorMode::Premultiplied` has been set, the alpha pixels must be premultiplied.
373    /// `AlphaColorMode::UnassociatedClean` has no effect in this function, and is equivalent to `AlphaColorMode::UnassociatedDirty`.
374    ///
375    /// returns AVIF file, size of color metadata, size of alpha metadata overhead
376    #[inline]
377    pub fn encode_raw_planes_10_bit(
378        &self, width: usize, height: usize,
379        planes: impl IntoIterator<Item = [u16; 3]> + Send,
380        alpha: Option<impl IntoIterator<Item = u16> + Send>,
381        color_pixel_range: PixelRange, matrix_coefficients: MatrixCoefficients,
382    ) -> Result<EncodedImage, Error> {
383        self.encode_raw_planes_internal(width, height, planes, alpha, color_pixel_range, matrix_coefficients, 10)
384    }
385
386    #[inline(never)]
387    fn encode_raw_planes_internal<P: rav1e::Pixel + Default>(
388        &self, width: usize, height: usize,
389        planes: impl IntoIterator<Item = [P; 3]> + Send,
390        alpha: Option<impl IntoIterator<Item = P> + Send>,
391        color_pixel_range: PixelRange, matrix_coefficients: MatrixCoefficients,
392        input_pixels_bit_depth: u8,
393    ) -> Result<EncodedImage, Error> {
394        let color_description = Some(ColorDescription {
395            transfer_characteristics: TransferCharacteristics::SRGB,
396            color_primaries: ColorPrimaries::BT709, // sRGB-compatible
397            matrix_coefficients,
398        });
399
400        let threads = self.threads.map(|threads| {
401            if threads > 0 { threads } else { rayon::current_num_threads() }
402        });
403
404        let encode_color = move || {
405            encode_to_av1::<P>(
406                &Av1EncodeConfig {
407                    width,
408                    height,
409                    bit_depth: input_pixels_bit_depth.into(),
410                    quantizer: self.quantizer.into(),
411                    speed: SpeedTweaks::from_my_preset(self.speed, self.quantizer),
412                    threads,
413                    pixel_range: color_pixel_range,
414                    chroma_sampling: ChromaSampling::Cs444,
415                    color_description,
416                },
417                move |frame| init_frame_3(width, height, planes, frame),
418            )
419        };
420        let encode_alpha = move || {
421            alpha.map(|alpha| {
422                encode_to_av1::<P>(
423                    &Av1EncodeConfig {
424                        width,
425                        height,
426                        bit_depth: input_pixels_bit_depth.into(),
427                        quantizer: self.alpha_quantizer.into(),
428                        speed: SpeedTweaks::from_my_preset(self.speed, self.alpha_quantizer),
429                        threads,
430                        pixel_range: PixelRange::Full,
431                        chroma_sampling: ChromaSampling::Cs400,
432                        color_description: None,
433                    },
434                    |frame| init_frame_1(width, height, alpha, frame),
435                )
436            })
437        };
438        #[cfg(all(target_arch = "wasm32", not(target_feature = "atomics")))]
439        let (color, alpha) = (encode_color(), encode_alpha());
440        #[cfg(not(all(target_arch = "wasm32", not(target_feature = "atomics"))))]
441        let (color, alpha) = rayon::join(encode_color, encode_alpha);
442        let (color, alpha) = (color?, alpha.transpose()?);
443
444        let mut serializer_config = avif_serialize::Aviffy::new();
445        serializer_config
446            .matrix_coefficients(match matrix_coefficients {
447                MatrixCoefficients::Identity => avif_serialize::constants::MatrixCoefficients::Rgb,
448                MatrixCoefficients::BT709 => avif_serialize::constants::MatrixCoefficients::Bt709,
449                MatrixCoefficients::Unspecified => avif_serialize::constants::MatrixCoefficients::Unspecified,
450                MatrixCoefficients::BT601 => avif_serialize::constants::MatrixCoefficients::Bt601,
451                MatrixCoefficients::YCgCo => avif_serialize::constants::MatrixCoefficients::Ycgco,
452                MatrixCoefficients::BT2020NCL => avif_serialize::constants::MatrixCoefficients::Bt2020Ncl,
453                MatrixCoefficients::BT2020CL => avif_serialize::constants::MatrixCoefficients::Bt2020Cl,
454                _ => return Err(Error::Unsupported("matrix coefficients")),
455            })
456            .premultiplied_alpha(self.premultiplied_alpha);
457        if let Some(exif) = &self.exif {
458            serializer_config.set_exif(exif.to_vec());
459        }
460        let avif_file = serializer_config.to_vec(&color, alpha.as_deref(), width as u32, height as u32, input_pixels_bit_depth);
461        let color_byte_size = color.len();
462        let alpha_byte_size = alpha.as_ref().map_or(0, |a| a.len());
463
464        Ok(EncodedImage {
465            avif_file, color_byte_size, alpha_byte_size,
466        })
467    }
468}
469
470/// Native endian
471#[inline(always)]
472fn to_ten(x: u8) -> u16 {
473    (u16::from(x) << 2) | (u16::from(x) >> 6)
474}
475
476/// Native endian
477#[inline(always)]
478fn rgb_to_10_bit_gbr(px: rgb::RGB<u8>) -> (u16, u16, u16) {
479    (to_ten(px.g), to_ten(px.b), to_ten(px.r))
480}
481
482#[inline(always)]
483fn rgb_to_8_bit_gbr(px: rgb::RGB<u8>) -> (u8, u8, u8) {
484    (px.g, px.b, px.r)
485}
486
487// const REC709: [f32; 3] = [0.2126, 0.7152, 0.0722];
488const BT601: [f32; 3] = [0.2990, 0.5870, 0.1140];
489
490#[inline(always)]
491fn rgb_to_ycbcr(px: rgb::RGB<u8>, depth: u8, matrix: [f32; 3]) -> (f32, f32, f32) {
492    let max_value = ((1 << depth) - 1) as f32;
493    let scale = max_value / 255.;
494    let shift = (max_value * 0.5).round();
495    let y = (scale * matrix[2]).mul_add(f32::from(px.b), (scale * matrix[0]).mul_add(f32::from(px.r), scale * matrix[1] * f32::from(px.g)));
496    let cb = f32::from(px.b).mul_add(scale, -y).mul_add(0.5 / (1. - matrix[2]), shift);
497    let cr = f32::from(px.r).mul_add(scale, -y).mul_add(0.5 / (1. - matrix[0]), shift);
498    (y.round(), cb.round(), cr.round())
499}
500
501#[inline(always)]
502fn rgb_to_10_bit_ycbcr(px: rgb::RGB<u8>, matrix: [f32; 3]) -> (u16, u16, u16) {
503    let (y, u, v) = rgb_to_ycbcr(px, 10, matrix);
504    (y as u16, u as u16, v as u16)
505}
506
507#[inline(always)]
508fn rgb_to_8_bit_ycbcr(px: rgb::RGB<u8>, matrix: [f32; 3]) -> (u8, u8, u8) {
509    let (y, u, v) = rgb_to_ycbcr(px, 8, matrix);
510    (y as u8, u as u8, v as u8)
511}
512
513fn quality_to_quantizer(quality: f32) -> u8 {
514    let q = quality / 100.;
515    let x = if q >= 0.82 { (1. - q) * 2.6 } else if q > 0.25 { q.mul_add(-0.5, 1. - 0.125) } else { 1. - q };
516    (x * 255.).round() as u8
517}
518
519#[derive(Debug, Copy, Clone)]
520struct SpeedTweaks {
521    pub speed_preset: u8,
522
523    pub fast_deblock: Option<bool>,
524    pub reduced_tx_set: Option<bool>,
525    pub tx_domain_distortion: Option<bool>,
526    pub tx_domain_rate: Option<bool>,
527    pub encode_bottomup: Option<bool>,
528    pub rdo_tx_decision: Option<bool>,
529    pub cdef: Option<bool>,
530    /// loop restoration filter
531    pub lrf: Option<bool>,
532    pub sgr_complexity_full: Option<bool>,
533    pub use_satd_subpel: Option<bool>,
534    pub inter_tx_split: Option<bool>,
535    pub fine_directional_intra: Option<bool>,
536    pub complex_prediction_modes: Option<bool>,
537    pub partition_range: Option<(u8, u8)>,
538    pub min_tile_size: u16,
539}
540
541impl SpeedTweaks {
542    pub fn from_my_preset(speed: u8, quantizer: u8) -> Self {
543        let low_quality = quantizer < quality_to_quantizer(55.);
544        let high_quality = quantizer > quality_to_quantizer(80.);
545        let max_block_size = if high_quality { 16 } else { 64 };
546
547        Self {
548            speed_preset: speed,
549
550            partition_range: Some(match speed {
551                0 => (4, 64.min(max_block_size)),
552                1 if low_quality => (4, 64.min(max_block_size)),
553                2 if low_quality => (4, 32.min(max_block_size)),
554                1..=4 => (4, 16),
555                5..=8 => (8, 16),
556                _ => (16, 16),
557            }),
558
559            complex_prediction_modes: Some(speed <= 1), // 2x-3x slower, 2% better
560            sgr_complexity_full: Some(speed <= 2), // 15% slower, barely improves anything -/+1%
561
562            encode_bottomup: Some(speed <= 2), // may be costly (+60%), may even backfire
563
564            // big blocks disabled at 3
565
566            // these two are together?
567            rdo_tx_decision: Some(speed <= 4 && !high_quality), // it tends to blur subtle textures
568            reduced_tx_set: Some(speed == 4 || speed >= 9), // It interacts with tx_domain_distortion too?
569
570            // 4px blocks disabled at 5
571
572            fine_directional_intra: Some(speed <= 6),
573            fast_deblock: Some(speed >= 7 && !high_quality), // mixed bag?
574
575            // 8px blocks disabled at 8
576            lrf: Some(low_quality && speed <= 8), // hardly any help for hi-q images. recovers some q at low quality
577            cdef: Some(low_quality && speed <= 9), // hardly any help for hi-q images. recovers some q at low quality
578
579            inter_tx_split: Some(speed >= 9), // mixed bag even when it works, and it backfires if not used together with reduced_tx_set
580            tx_domain_rate: Some(speed >= 10), // 20% faster, but also 10% larger files!
581
582            tx_domain_distortion: None, // very mixed bag, sometimes helps speed sometimes it doesn't
583            use_satd_subpel: Some(false), // doesn't make sense
584            min_tile_size: match speed {
585                0 => 4096,
586                1 => 2048,
587                2 => 1024,
588                3 => 512,
589                4 => 256,
590                _ => 128,
591            } * if high_quality { 2 } else { 1 },
592        }
593    }
594
595    pub(crate) fn speed_settings(&self) -> SpeedSettings {
596        let mut speed_settings = SpeedSettings::from_preset(self.speed_preset);
597
598        speed_settings.multiref = false;
599        speed_settings.rdo_lookahead_frames = 1;
600        speed_settings.scene_detection_mode = SceneDetectionSpeed::None;
601        speed_settings.motion.include_near_mvs = false;
602
603        if let Some(v) = self.fast_deblock { speed_settings.fast_deblock = v; }
604        if let Some(v) = self.reduced_tx_set { speed_settings.transform.reduced_tx_set = v; }
605        if let Some(v) = self.tx_domain_distortion { speed_settings.transform.tx_domain_distortion = v; }
606        if let Some(v) = self.tx_domain_rate { speed_settings.transform.tx_domain_rate = v; }
607        if let Some(v) = self.encode_bottomup { speed_settings.partition.encode_bottomup = v; }
608        if let Some(v) = self.rdo_tx_decision { speed_settings.transform.rdo_tx_decision = v; }
609        if let Some(v) = self.cdef { speed_settings.cdef = v; }
610        if let Some(v) = self.lrf { speed_settings.lrf = v; }
611        if let Some(v) = self.inter_tx_split { speed_settings.transform.enable_inter_tx_split = v; }
612        if let Some(v) = self.sgr_complexity_full { speed_settings.sgr_complexity = if v { SGRComplexityLevel::Full } else { SGRComplexityLevel::Reduced } }
613        if let Some(v) = self.use_satd_subpel { speed_settings.motion.use_satd_subpel = v; }
614        if let Some(v) = self.fine_directional_intra { speed_settings.prediction.fine_directional_intra = v; }
615        if let Some(v) = self.complex_prediction_modes { speed_settings.prediction.prediction_modes = if v { PredictionModesSetting::ComplexAll } else { PredictionModesSetting::Simple} }
616        if let Some((min, max)) = self.partition_range {
617            debug_assert!(min <= max);
618            fn sz(s: u8) -> BlockSize {
619                match s {
620                    4 => BlockSize::BLOCK_4X4,
621                    8 => BlockSize::BLOCK_8X8,
622                    16 => BlockSize::BLOCK_16X16,
623                    32 => BlockSize::BLOCK_32X32,
624                    64 => BlockSize::BLOCK_64X64,
625                    128 => BlockSize::BLOCK_128X128,
626                    _ => panic!("bad size {s}"),
627                }
628            }
629            speed_settings.partition.partition_range = PartitionRange::new(sz(min), sz(max));
630        }
631
632        speed_settings
633    }
634}
635
636struct Av1EncodeConfig {
637    pub width: usize,
638    pub height: usize,
639    pub bit_depth: usize,
640    pub quantizer: usize,
641    pub speed: SpeedTweaks,
642    /// 0 means num_cpus
643    pub threads: Option<usize>,
644    pub pixel_range: PixelRange,
645    pub chroma_sampling: ChromaSampling,
646    pub color_description: Option<ColorDescription>,
647}
648
649fn rav1e_config(p: &Av1EncodeConfig) -> Config {
650    // AV1 needs all the CPU power you can give it,
651    // except when it'd create inefficiently tiny tiles
652    let tiles = {
653        let threads = p.threads.unwrap_or_else(rayon::current_num_threads);
654        threads.min((p.width * p.height) / (p.speed.min_tile_size as usize).pow(2))
655    };
656    let speed_settings = p.speed.speed_settings();
657    let cfg = Config::new()
658        .with_encoder_config(EncoderConfig {
659        width: p.width,
660        height: p.height,
661        time_base: Rational::new(1, 1),
662        sample_aspect_ratio: Rational::new(1, 1),
663        bit_depth: p.bit_depth,
664        chroma_sampling: p.chroma_sampling,
665        chroma_sample_position: ChromaSamplePosition::Unknown,
666        pixel_range: p.pixel_range,
667        color_description: p.color_description,
668        mastering_display: None,
669        content_light: None,
670        enable_timing_info: false,
671        still_picture: true,
672        error_resilient: false,
673        switch_frame_interval: 0,
674        min_key_frame_interval: 0,
675        max_key_frame_interval: 0,
676        reservoir_frame_delay: None,
677        low_latency: false,
678        quantizer: p.quantizer,
679        min_quantizer: p.quantizer as _,
680        bitrate: 0,
681        tune: Tune::Psychovisual,
682        tile_cols: 0,
683        tile_rows: 0,
684        tiles,
685        film_grain_params: None,
686        level_idx: None,
687        speed_settings,
688    });
689
690    if let Some(threads) = p.threads {
691        cfg.with_threads(threads)
692    } else {
693        cfg
694    }
695}
696
697fn init_frame_3<P: rav1e::Pixel + Default>(
698    width: usize, height: usize, planes: impl IntoIterator<Item = [P; 3]> + Send, frame: &mut Frame<P>,
699) -> Result<(), Error> {
700    let mut f = frame.planes.iter_mut();
701    let mut planes = planes.into_iter();
702
703    // it doesn't seem to be necessary to fill padding area
704    let mut y = f.next().unwrap().mut_slice(Default::default());
705    let mut u = f.next().unwrap().mut_slice(Default::default());
706    let mut v = f.next().unwrap().mut_slice(Default::default());
707
708    for ((y, u), v) in y.rows_iter_mut().zip(u.rows_iter_mut()).zip(v.rows_iter_mut()).take(height) {
709        let y = &mut y[..width];
710        let u = &mut u[..width];
711        let v = &mut v[..width];
712        for ((y, u), v) in y.iter_mut().zip(u).zip(v) {
713            let px = planes.next().ok_or(Error::TooFewPixels)?;
714            *y = px[0];
715            *u = px[1];
716            *v = px[2];
717        }
718    }
719    Ok(())
720}
721
722fn init_frame_1<P: rav1e::Pixel + Default>(width: usize, height: usize, planes: impl IntoIterator<Item = P> + Send, frame: &mut Frame<P>) -> Result<(), Error> {
723    let mut y = frame.planes[0].mut_slice(Default::default());
724    let mut planes = planes.into_iter();
725
726    for y in y.rows_iter_mut().take(height) {
727        let y = &mut y[..width];
728        for y in y.iter_mut() {
729            *y = planes.next().ok_or(Error::TooFewPixels)?;
730        }
731    }
732    Ok(())
733}
734
735#[inline(never)]
736fn encode_to_av1<P: rav1e::Pixel>(p: &Av1EncodeConfig, init: impl FnOnce(&mut Frame<P>) -> Result<(), Error>) -> Result<Vec<u8>, Error> {
737    let mut ctx: Context<P> = rav1e_config(p).new_context()?;
738    let mut frame = ctx.new_frame();
739
740    init(&mut frame)?;
741    ctx.send_frame(frame)?;
742    ctx.flush();
743
744    let mut out = Vec::new();
745    loop {
746        match ctx.receive_packet() {
747            Ok(mut packet) => match packet.frame_type {
748                FrameType::KEY => {
749                    out.append(&mut packet.data);
750                },
751                _ => continue,
752            },
753            Err(EncoderStatus::Encoded | EncoderStatus::LimitReached) => break,
754            Err(err) => Err(err)?,
755        }
756    }
757    Ok(out)
758}
ravif/av1encoder.rs

ravif/
av1encoder.rs
