1 | #![allow(deprecated)] |
---|---|
2 | use crate::dirtyalpha::blurred_dirty_alpha; |
3 | use crate::error::Error; |
4 | #[cfg(not(feature = "threading"))] |
5 | use crate::rayoff as rayon; |
6 | use imgref::{Img, ImgVec}; |
7 | use rav1e::prelude::*; |
8 | use rgb::{RGB8, RGBA8}; |
9 | |
10 | /// For [`Encoder::with_internal_color_model`] |
11 | #[derive(Debug, Copy, Clone, Eq, PartialEq)] |
12 | pub 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)] |
25 | pub 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)] |
42 | pub enum BitDepth { |
43 | Eight, |
44 | Ten, |
45 | /// Pick 8 or 10 depending on image format and decoder compatibility |
46 | #[default] |
47 | Auto, |
48 | } |
49 | |
50 | /// The newly-created image file + extra info FYI |
51 | #[non_exhaustive] |
52 | #[derive(Clone)] |
53 | pub 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)] |
64 | pub 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 |
84 | impl 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 |
193 | impl 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 | BitDepth::Auto => { |
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 => { |
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 | .filter(|px| px.a != 255) |
261 | .map(|px| { |
262 | if px.a == 0 { |
263 | RGBA8::default() |
264 | } else { |
265 | RGBA8::new( |
266 | (u16::from(px.r) * 255 / u16::from(px.a)) as u8, |
267 | (u16::from(px.g) * 255 / u16::from(px.a)) as u8, |
268 | (u16::from(px.b) * 255 / u16::from(px.a)) as u8, |
269 | px.a, |
270 | ) |
271 | } |
272 | }) |
273 | .collect(); |
274 | Some(ImgVec::new(prem, in_buffer.width(), in_buffer.height())) |
275 | }, |
276 | } |
277 | } |
278 | |
279 | /// Make a new AVIF image from RGB pixels |
280 | /// |
281 | /// Make the `Img` for the `buffer` like this: |
282 | /// |
283 | /// ```rust,ignore |
284 | /// Img::new(&pixels_rgb[..], width, height) |
285 | /// ``` |
286 | /// |
287 | /// If you have pixels as `u8` slice, then first do: |
288 | /// |
289 | /// ```rust,ignore |
290 | /// use rgb::ComponentSlice; |
291 | /// let pixels_rgb = pixels_u8.as_rgb(); |
292 | /// ``` |
293 | /// |
294 | /// returns AVIF file, size of color metadata |
295 | #[inline] |
296 | pub fn encode_rgb(&self, buffer: Img<&[RGB8]>) -> Result<EncodedImage, Error> { |
297 | self.encode_rgb_internal_from_8bit(buffer.width(), buffer.height(), buffer.pixels()) |
298 | } |
299 | |
300 | fn encode_rgb_internal_from_8bit(&self, width: usize, height: usize, pixels: impl Iterator<Item = RGB8> + Send + Sync) -> Result<EncodedImage, Error> { |
301 | let matrix_coefficients = match self.color_model { |
302 | ColorModel::YCbCr => MatrixCoefficients::BT601, |
303 | ColorModel::RGB => MatrixCoefficients::Identity, |
304 | }; |
305 | |
306 | match self.output_depth { |
307 | BitDepth::Eight => { |
308 | let planes = pixels.map(|px| { |
309 | let (y, u, v) = match self.color_model { |
310 | ColorModel::YCbCr => rgb_to_8_bit_ycbcr(px, BT601), |
311 | ColorModel::RGB => rgb_to_8_bit_gbr(px), |
312 | }; |
313 | [y, u, v] |
314 | }); |
315 | self.encode_raw_planes_8_bit(width, height, planes, None::<[_; 0]>, PixelRange::Full, matrix_coefficients) |
316 | }, |
317 | BitDepth::Ten | BitDepth::Auto => { |
318 | let planes = pixels.map(|px| { |
319 | let (y, u, v) = match self.color_model { |
320 | ColorModel::YCbCr => rgb_to_10_bit_ycbcr(px, BT601), |
321 | ColorModel::RGB => rgb_to_10_bit_gbr(px), |
322 | }; |
323 | [y, u, v] |
324 | }); |
325 | self.encode_raw_planes_10_bit(width, height, planes, None::<[_; 0]>, PixelRange::Full, matrix_coefficients) |
326 | }, |
327 | } |
328 | } |
329 | |
330 | /// Encodes AVIF from 3 planar channels that are in the color space described by `matrix_coefficients`, |
331 | /// with sRGB transfer characteristics and color primaries. |
332 | /// |
333 | /// Alpha always uses full range. Chroma subsampling is not supported, and it's a bad idea for AVIF anyway. |
334 | /// If there's no alpha, use `None::<[_; 0]>`. |
335 | /// |
336 | /// `color_pixel_range` should be `PixelRange::Full` to avoid worsening already small 8-bit dynamic range. |
337 | /// Support for limited range may be removed in the future. |
338 | /// |
339 | /// If `AlphaColorMode::Premultiplied` has been set, the alpha pixels must be premultiplied. |
340 | /// `AlphaColorMode::UnassociatedClean` has no effect in this function, and is equivalent to `AlphaColorMode::UnassociatedDirty`. |
341 | /// |
342 | /// returns AVIF file, size of color metadata, size of alpha metadata overhead |
343 | #[inline] |
344 | pub fn encode_raw_planes_8_bit( |
345 | &self, width: usize, height: usize, planes: impl IntoIterator<Item = [u8; 3]> + Send, alpha: Option<impl IntoIterator<Item = u8> + Send>, |
346 | color_pixel_range: PixelRange, matrix_coefficients: MatrixCoefficients, |
347 | ) -> Result<EncodedImage, Error> { |
348 | self.encode_raw_planes_internal(width, height, planes, alpha, color_pixel_range, matrix_coefficients, 8) |
349 | } |
350 | |
351 | /// Encodes AVIF from 3 planar channels that are in the color space described by `matrix_coefficients`, |
352 | /// with sRGB transfer characteristics and color primaries. |
353 | /// |
354 | /// The pixels are 10-bit (values `0.=1023`). |
355 | /// |
356 | /// Alpha always uses full range. Chroma subsampling is not supported, and it's a bad idea for AVIF anyway. |
357 | /// If there's no alpha, use `None::<[_; 0]>`. |
358 | /// |
359 | /// `color_pixel_range` should be `PixelRange::Full`. Support for limited range may be removed in the future. |
360 | /// |
361 | /// If `AlphaColorMode::Premultiplied` has been set, the alpha pixels must be premultiplied. |
362 | /// `AlphaColorMode::UnassociatedClean` has no effect in this function, and is equivalent to `AlphaColorMode::UnassociatedDirty`. |
363 | /// |
364 | /// returns AVIF file, size of color metadata, size of alpha metadata overhead |
365 | #[inline] |
366 | pub fn encode_raw_planes_10_bit( |
367 | &self, width: usize, height: usize, planes: impl IntoIterator<Item = [u16; 3]> + Send, alpha: Option<impl IntoIterator<Item = u16> + Send>, |
368 | color_pixel_range: PixelRange, matrix_coefficients: MatrixCoefficients, |
369 | ) -> Result<EncodedImage, Error> { |
370 | self.encode_raw_planes_internal(width, height, planes, alpha, color_pixel_range, matrix_coefficients, 10) |
371 | } |
372 | |
373 | #[inline(never)] |
374 | fn encode_raw_planes_internal<P: rav1e::Pixel + Default>( |
375 | &self, width: usize, height: usize, planes: impl IntoIterator<Item = [P; 3]> + Send, alpha: Option<impl IntoIterator<Item = P> + Send>, |
376 | color_pixel_range: PixelRange, matrix_coefficients: MatrixCoefficients, input_pixels_bit_depth: u8, |
377 | ) -> Result<EncodedImage, Error> { |
378 | let color_description = Some(ColorDescription { |
379 | transfer_characteristics: TransferCharacteristics::SRGB, |
380 | color_primaries: ColorPrimaries::BT709, // sRGB-compatible |
381 | matrix_coefficients, |
382 | }); |
383 | |
384 | let threads = self.threads.map(|threads| { |
385 | if threads > 0 { threads } else { rayon::current_num_threads() } |
386 | }); |
387 | |
388 | let encode_color = move || { |
389 | encode_to_av1::<P>( |
390 | &Av1EncodeConfig { |
391 | width, |
392 | height, |
393 | bit_depth: input_pixels_bit_depth.into(), |
394 | quantizer: self.quantizer.into(), |
395 | speed: SpeedTweaks::from_my_preset(self.speed, self.quantizer), |
396 | threads, |
397 | pixel_range: color_pixel_range, |
398 | chroma_sampling: ChromaSampling::Cs444, |
399 | color_description, |
400 | }, |
401 | move |frame| init_frame_3(width, height, planes, frame), |
402 | ) |
403 | }; |
404 | let encode_alpha = move || { |
405 | alpha.map(|alpha| { |
406 | encode_to_av1::<P>( |
407 | &Av1EncodeConfig { |
408 | width, |
409 | height, |
410 | bit_depth: input_pixels_bit_depth.into(), |
411 | quantizer: self.alpha_quantizer.into(), |
412 | speed: SpeedTweaks::from_my_preset(self.speed, self.alpha_quantizer), |
413 | threads, |
414 | pixel_range: PixelRange::Full, |
415 | chroma_sampling: ChromaSampling::Cs400, |
416 | color_description: None, |
417 | }, |
418 | |frame| init_frame_1(width, height, alpha, frame), |
419 | ) |
420 | }) |
421 | }; |
422 | #[cfg(all(target_arch = "wasm32", not(target_feature = "atomics")))] |
423 | let (color, alpha) = (encode_color(), encode_alpha()); |
424 | #[cfg(not(all(target_arch = "wasm32", not(target_feature = "atomics"))))] |
425 | let (color, alpha) = rayon::join(encode_color, encode_alpha); |
426 | let (color, alpha) = (color?, alpha.transpose()?); |
427 | |
428 | let avif_file = avif_serialize::Aviffy::new() |
429 | .matrix_coefficients(match matrix_coefficients { |
430 | MatrixCoefficients::Identity => avif_serialize::constants::MatrixCoefficients::Rgb, |
431 | MatrixCoefficients::BT709 => avif_serialize::constants::MatrixCoefficients::Bt709, |
432 | MatrixCoefficients::Unspecified => avif_serialize::constants::MatrixCoefficients::Unspecified, |
433 | MatrixCoefficients::BT601 => avif_serialize::constants::MatrixCoefficients::Bt601, |
434 | MatrixCoefficients::YCgCo => avif_serialize::constants::MatrixCoefficients::Ycgco, |
435 | MatrixCoefficients::BT2020NCL => avif_serialize::constants::MatrixCoefficients::Bt2020Ncl, |
436 | MatrixCoefficients::BT2020CL => avif_serialize::constants::MatrixCoefficients::Bt2020Cl, |
437 | _ => return Err(Error::Unsupported("matrix coefficients")), |
438 | }) |
439 | .premultiplied_alpha(self.premultiplied_alpha) |
440 | .to_vec(&color, alpha.as_deref(), width as u32, height as u32, input_pixels_bit_depth); |
441 | let color_byte_size = color.len(); |
442 | let alpha_byte_size = alpha.as_ref().map_or(0, |a| a.len()); |
443 | |
444 | Ok(EncodedImage { |
445 | avif_file, color_byte_size, alpha_byte_size, |
446 | }) |
447 | } |
448 | } |
449 | |
450 | #[inline(always)] |
451 | fn to_ten(x: u8) -> u16 { |
452 | (u16::from(x) << 2) | (u16::from(x) >> 6) |
453 | } |
454 | |
455 | #[inline(always)] |
456 | fn rgb_to_10_bit_gbr(px: rgb::RGB<u8>) -> (u16, u16, u16) { |
457 | (to_ten(px.g), to_ten(px.b), to_ten(px.r)) |
458 | } |
459 | |
460 | #[inline(always)] |
461 | fn rgb_to_8_bit_gbr(px: rgb::RGB<u8>) -> (u8, u8, u8) { |
462 | (px.g, px.b, px.r) |
463 | } |
464 | |
465 | // const REC709: [f32; 3] = [0.2126, 0.7152, 0.0722]; |
466 | const BT601: [f32; 3] = [0.2990, 0.5870, 0.1140]; |
467 | |
468 | #[inline(always)] |
469 | fn rgb_to_ycbcr(px: rgb::RGB<u8>, depth: u8, matrix: [f32; 3]) -> (f32, f32, f32) { |
470 | let max_value: f32 = ((1 << depth) - 1) as f32; |
471 | let scale: f32 = max_value / 255.; |
472 | let shift: f32 = (max_value * 0.5).round(); |
473 | let y: f32 = scale * matrix[0] * f32::from(px.r) + scale * matrix[1] * f32::from(px.g) + scale * matrix[2] * f32::from(px.b); |
474 | let cb: f32 = (f32::from(px.b) * scale - y).mul_add(a:0.5 / (1. - matrix[2]), b:shift); |
475 | let cr: f32 = (f32::from(px.r) * scale - y).mul_add(a:0.5 / (1. - matrix[0]), b:shift); |
476 | (y.round(), cb.round(), cr.round()) |
477 | } |
478 | |
479 | #[inline(always)] |
480 | fn rgb_to_10_bit_ycbcr(px: rgb::RGB<u8>, matrix: [f32; 3]) -> (u16, u16, u16) { |
481 | let (y: f32, u: f32, v: f32) = rgb_to_ycbcr(px, depth:10, matrix); |
482 | (y as u16, u as u16, v as u16) |
483 | } |
484 | |
485 | #[inline(always)] |
486 | fn rgb_to_8_bit_ycbcr(px: rgb::RGB<u8>, matrix: [f32; 3]) -> (u8, u8, u8) { |
487 | let (y: f32, u: f32, v: f32) = rgb_to_ycbcr(px, depth:8, matrix); |
488 | (y as u8, u as u8, v as u8) |
489 | } |
490 | |
491 | fn quality_to_quantizer(quality: f32) -> u8 { |
492 | let q: f32 = quality / 100.; |
493 | let x: f32 = if q >= 0.85 { (1. - q) * 3. } else if q > 0.25 { 1. - 0.125 - q * 0.5 } else { 1. - q }; |
494 | (x * 255.).round() as u8 |
495 | } |
496 | |
497 | #[derive(Debug, Copy, Clone)] |
498 | struct SpeedTweaks { |
499 | pub speed_preset: u8, |
500 | |
501 | pub fast_deblock: Option<bool>, |
502 | pub reduced_tx_set: Option<bool>, |
503 | pub tx_domain_distortion: Option<bool>, |
504 | pub tx_domain_rate: Option<bool>, |
505 | pub encode_bottomup: Option<bool>, |
506 | pub rdo_tx_decision: Option<bool>, |
507 | pub cdef: Option<bool>, |
508 | /// loop restoration filter |
509 | pub lrf: Option<bool>, |
510 | pub sgr_complexity_full: Option<bool>, |
511 | pub use_satd_subpel: Option<bool>, |
512 | pub inter_tx_split: Option<bool>, |
513 | pub fine_directional_intra: Option<bool>, |
514 | pub complex_prediction_modes: Option<bool>, |
515 | pub partition_range: Option<(u8, u8)>, |
516 | pub min_tile_size: u16, |
517 | } |
518 | |
519 | impl SpeedTweaks { |
520 | pub fn from_my_preset(speed: u8, quantizer: u8) -> Self { |
521 | let low_quality = quantizer < quality_to_quantizer(55.); |
522 | let high_quality = quantizer > quality_to_quantizer(80.); |
523 | let max_block_size = if high_quality { 16 } else { 64 }; |
524 | |
525 | Self { |
526 | speed_preset: speed, |
527 | |
528 | partition_range: Some(match speed { |
529 | 0 => (4, 64.min(max_block_size)), |
530 | 1 if low_quality => (4, 64.min(max_block_size)), |
531 | 2 if low_quality => (4, 32.min(max_block_size)), |
532 | 1..=4 => (4, 16), |
533 | 5..=8 => (8, 16), |
534 | _ => (16, 16), |
535 | }), |
536 | |
537 | complex_prediction_modes: Some(speed <= 1), // 2x-3x slower, 2% better |
538 | sgr_complexity_full: Some(speed <= 2), // 15% slower, barely improves anything -/+1% |
539 | |
540 | encode_bottomup: Some(speed <= 2), // may be costly (+60%), may even backfire |
541 | |
542 | // big blocks disabled at 3 |
543 | |
544 | // these two are together? |
545 | rdo_tx_decision: Some(speed <= 4 && !high_quality), // it tends to blur subtle textures |
546 | reduced_tx_set: Some(speed == 4 || speed >= 9), // It interacts with tx_domain_distortion too? |
547 | |
548 | // 4px blocks disabled at 5 |
549 | |
550 | fine_directional_intra: Some(speed <= 6), |
551 | fast_deblock: Some(speed >= 7 && !high_quality), // mixed bag? |
552 | |
553 | // 8px blocks disabled at 8 |
554 | lrf: Some(low_quality && speed <= 8), // hardly any help for hi-q images. recovers some q at low quality |
555 | cdef: Some(low_quality && speed <= 9), // hardly any help for hi-q images. recovers some q at low quality |
556 | |
557 | inter_tx_split: Some(speed >= 9), // mixed bag even when it works, and it backfires if not used together with reduced_tx_set |
558 | tx_domain_rate: Some(speed >= 10), // 20% faster, but also 10% larger files! |
559 | |
560 | tx_domain_distortion: None, // very mixed bag, sometimes helps speed sometimes it doesn't |
561 | use_satd_subpel: Some(false), // doesn't make sense |
562 | min_tile_size: match speed { |
563 | 0 => 4096, |
564 | 1 => 2048, |
565 | 2 => 1024, |
566 | 3 => 512, |
567 | 4 => 256, |
568 | _ => 128, |
569 | } * if high_quality { 2 } else { 1 }, |
570 | } |
571 | } |
572 | |
573 | pub(crate) fn speed_settings(&self) -> SpeedSettings { |
574 | let mut speed_settings = SpeedSettings::from_preset(self.speed_preset); |
575 | |
576 | speed_settings.multiref = false; |
577 | speed_settings.rdo_lookahead_frames = 1; |
578 | speed_settings.scene_detection_mode = SceneDetectionSpeed::None; |
579 | speed_settings.motion.include_near_mvs = false; |
580 | |
581 | if let Some(v) = self.fast_deblock { speed_settings.fast_deblock = v; } |
582 | if let Some(v) = self.reduced_tx_set { speed_settings.transform.reduced_tx_set = v; } |
583 | if let Some(v) = self.tx_domain_distortion { speed_settings.transform.tx_domain_distortion = v; } |
584 | if let Some(v) = self.tx_domain_rate { speed_settings.transform.tx_domain_rate = v; } |
585 | if let Some(v) = self.encode_bottomup { speed_settings.partition.encode_bottomup = v; } |
586 | if let Some(v) = self.rdo_tx_decision { speed_settings.transform.rdo_tx_decision = v; } |
587 | if let Some(v) = self.cdef { speed_settings.cdef = v; } |
588 | if let Some(v) = self.lrf { speed_settings.lrf = v; } |
589 | if let Some(v) = self.inter_tx_split { speed_settings.transform.enable_inter_tx_split = v; } |
590 | if let Some(v) = self.sgr_complexity_full { speed_settings.sgr_complexity = if v { SGRComplexityLevel::Full } else { SGRComplexityLevel::Reduced } }; |
591 | if let Some(v) = self.use_satd_subpel { speed_settings.motion.use_satd_subpel = v; } |
592 | if let Some(v) = self.fine_directional_intra { speed_settings.prediction.fine_directional_intra = v; } |
593 | if let Some(v) = self.complex_prediction_modes { speed_settings.prediction.prediction_modes = if v { PredictionModesSetting::ComplexAll } else { PredictionModesSetting::Simple} }; |
594 | if let Some((min, max)) = self.partition_range { |
595 | debug_assert!(min <= max); |
596 | fn sz(s: u8) -> BlockSize { |
597 | match s { |
598 | 4 => BlockSize::BLOCK_4X4, |
599 | 8 => BlockSize::BLOCK_8X8, |
600 | 16 => BlockSize::BLOCK_16X16, |
601 | 32 => BlockSize::BLOCK_32X32, |
602 | 64 => BlockSize::BLOCK_64X64, |
603 | 128 => BlockSize::BLOCK_128X128, |
604 | _ => panic!("bad size{s} "), |
605 | } |
606 | } |
607 | speed_settings.partition.partition_range = PartitionRange::new(sz(min), sz(max)); |
608 | } |
609 | |
610 | speed_settings |
611 | } |
612 | } |
613 | |
614 | struct Av1EncodeConfig { |
615 | pub width: usize, |
616 | pub height: usize, |
617 | pub bit_depth: usize, |
618 | pub quantizer: usize, |
619 | pub speed: SpeedTweaks, |
620 | /// 0 means num_cpus |
621 | pub threads: Option<usize>, |
622 | pub pixel_range: PixelRange, |
623 | pub chroma_sampling: ChromaSampling, |
624 | pub color_description: Option<ColorDescription>, |
625 | } |
626 | |
627 | fn rav1e_config(p: &Av1EncodeConfig) -> Config { |
628 | // AV1 needs all the CPU power you can give it, |
629 | // except when it'd create inefficiently tiny tiles |
630 | let tiles = { |
631 | let threads = p.threads.unwrap_or_else(rayon::current_num_threads); |
632 | threads.min((p.width * p.height) / (p.speed.min_tile_size as usize).pow(2)) |
633 | }; |
634 | let speed_settings = p.speed.speed_settings(); |
635 | let cfg = Config::new() |
636 | .with_encoder_config(EncoderConfig { |
637 | width: p.width, |
638 | height: p.height, |
639 | time_base: Rational::new(1, 1), |
640 | sample_aspect_ratio: Rational::new(1, 1), |
641 | bit_depth: p.bit_depth, |
642 | chroma_sampling: p.chroma_sampling, |
643 | chroma_sample_position: ChromaSamplePosition::Unknown, |
644 | pixel_range: p.pixel_range, |
645 | color_description: p.color_description, |
646 | mastering_display: None, |
647 | content_light: None, |
648 | enable_timing_info: false, |
649 | still_picture: true, |
650 | error_resilient: false, |
651 | switch_frame_interval: 0, |
652 | min_key_frame_interval: 0, |
653 | max_key_frame_interval: 0, |
654 | reservoir_frame_delay: None, |
655 | low_latency: false, |
656 | quantizer: p.quantizer, |
657 | min_quantizer: p.quantizer as _, |
658 | bitrate: 0, |
659 | tune: Tune::Psychovisual, |
660 | tile_cols: 0, |
661 | tile_rows: 0, |
662 | tiles, |
663 | film_grain_params: None, |
664 | level_idx: None, |
665 | speed_settings, |
666 | }); |
667 | |
668 | if let Some(threads) = p.threads { |
669 | cfg.with_threads(threads) |
670 | } else { |
671 | cfg |
672 | } |
673 | } |
674 | |
675 | fn init_frame_3<P: rav1e::Pixel + Default>( |
676 | width: usize, height: usize, planes: impl IntoIterator<Item = [P; 3]> + Send, frame: &mut Frame<P>, |
677 | ) -> Result<(), Error> { |
678 | let mut f: IterMut<'_, Plane > = frame.planes.iter_mut(); |
679 | let mut planes: |
680 | |
681 | // it doesn't seem to be necessary to fill padding area |
682 | let mut y: PlaneMutSlice<'_, P> = f.next().unwrap().mut_slice(po:Default::default()); |
683 | let mut u: PlaneMutSlice<'_, P> = f.next().unwrap().mut_slice(po:Default::default()); |
684 | let mut v: PlaneMutSlice<'_, P> = f.next().unwrap().mut_slice(po:Default::default()); |
685 | |
686 | for ((y: &mut [P], u: &mut [P]), v: &mut [P]) in y.rows_iter_mut().zip(u.rows_iter_mut()).zip(v.rows_iter_mut()).take(height) { |
687 | let y: &mut [P] = &mut y[..width]; |
688 | let u: &mut [P] = &mut u[..width]; |
689 | let v: &mut [P] = &mut v[..width]; |
690 | for ((y: &mut P, u: &mut P), v: &mut P) in y.iter_mut().zip(u).zip(v) { |
691 | let px: [P; 3] = planes.next().ok_or(err:Error::TooFewPixels)?; |
692 | *y = px[0]; |
693 | *u = px[1]; |
694 | *v = px[2]; |
695 | } |
696 | } |
697 | Ok(()) |
698 | } |
699 | |
700 | fn init_frame_1<P: rav1e::Pixel + Default>(width: usize, height: usize, planes: impl IntoIterator<Item = P> + Send, frame: &mut Frame<P>) -> Result<(), Error> { |
701 | let mut y: PlaneMutSlice<'_, P> = frame.planes[0].mut_slice(po:Default::default()); |
702 | let mut planes: |
703 | |
704 | for y: &mut [P] in y.rows_iter_mut().take(height) { |
705 | let y: &mut [P] = &mut y[..width]; |
706 | for y: &mut P in y.iter_mut() { |
707 | *y = planes.next().ok_or(err:Error::TooFewPixels)?; |
708 | } |
709 | } |
710 | Ok(()) |
711 | } |
712 | |
713 | #[inline(never)] |
714 | fn encode_to_av1<P: rav1e::Pixel>(p: &Av1EncodeConfig, init: impl FnOnce(&mut Frame<P>) -> Result<(), Error>) -> Result<Vec<u8>, Error> { |
715 | let mut ctx: Context<P> = rav1e_config(p).new_context()?; |
716 | let mut frame: Frame = ctx.new_frame(); |
717 | |
718 | init(&mut frame)?; |
719 | ctx.send_frame(frame)?; |
720 | ctx.flush(); |
721 | |
722 | let mut out: Vec |
723 | loop { |
724 | match ctx.receive_packet() { |
725 | Ok(mut packet: Packet) => match packet.frame_type { |
726 | FrameType::KEY => { |
727 | out.append(&mut packet.data); |
728 | }, |
729 | _ => continue, |
730 | }, |
731 | Err(EncoderStatus::Encoded) | |
732 | Err(EncoderStatus::LimitReached) => break, |
733 | Err(err: EncoderStatus) => Err(err)?, |
734 | } |
735 | } |
736 | Ok(out) |
737 | } |
738 |