ravif/
av1encoder.rs

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

ravif/
av1encoder.rs
