ravif/
av1encoder.rs

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

ravif/
av1encoder.rs
