webp_rust/encoder/lossy/
api.rs

1//! Public entry points and final frame assembly for lossy encoding.
2
3use super::bitstream::*;
4use super::predict::*;
5use super::*;
6use crate::encoder::writer::ByteWriter;
7
8/// Builds a raw VP8 frame from the already encoded partitions.
9fn build_vp8_frame(
10    width: usize,
11    height: usize,
12    partition0: &[u8],
13    token_partition: &[u8],
14) -> Result<Vec<u8>, EncoderError> {
15    if partition0.len() > MAX_PARTITION0_LENGTH {
16        return Err(EncoderError::Bitstream("VP8 partition 0 overflow"));
17    }
18
19    let payload_size = 10usize
20        .checked_add(partition0.len())
21        .and_then(|size| size.checked_add(token_partition.len()))
22        .ok_or(EncoderError::InvalidParam("encoded output is too large"))?;
23
24    let mut data = ByteWriter::with_capacity(payload_size);
25    let frame_bits = ((partition0.len() as u32) << 5) | (1 << 4);
26    data.write_u24_le(frame_bits);
27    data.write_bytes(&[0x9d, 0x01, 0x2a]);
28    data.write_u16_le(width as u16);
29    data.write_u16_le(height as u16);
30    data.write_bytes(partition0);
31    data.write_bytes(token_partition);
32    Ok(data.into_bytes())
33}
34
35/// Builds a VP8 frame for a candidate mode/filter combination.
36fn build_candidate_vp8_frame(
37    width: usize,
38    height: usize,
39    mb_width: usize,
40    mb_height: usize,
41    candidate: &EncodedLossyCandidate,
42    filter: &FilterConfig,
43) -> Result<Vec<u8>, EncoderError> {
44    let segment = segment_with_uniform_filter(&candidate.segment, filter.level);
45    let partition0 = encode_partition0(
46        mb_width,
47        mb_height,
48        candidate.base_quant,
49        &segment,
50        filter,
51        &candidate.probabilities,
52        &candidate.modes,
53    );
54    build_vp8_frame(width, height, &partition0, &candidate.token_partition)
55}
56
57/// Encodes one lossy candidate and captures its token partition, probabilities, and modes.
58fn encode_lossy_candidate(
59    source: &Planes,
60    mb_width: usize,
61    mb_height: usize,
62    profile: &LossySearchProfile,
63    segment: &SegmentConfig,
64) -> Result<EncodedLossyCandidate, EncoderError> {
65    let segment_quants = build_segment_quantizers(segment);
66    let (token_partition, probabilities, modes) = if profile.update_probabilities {
67        let mut stats = [[[[0u32; NUM_PROBAS]; NUM_CTX]; NUM_BANDS]; NUM_TYPES];
68        let (initial_partition, _, initial_modes) = encode_token_partition(
69            source,
70            mb_width,
71            mb_height,
72            profile,
73            segment,
74            &segment_quants,
75            &COEFFS_PROBA0,
76            Some(&mut stats),
77        );
78        let probabilities = finalize_token_probabilities(&stats);
79        if probabilities == COEFFS_PROBA0 {
80            (initial_partition, probabilities, initial_modes)
81        } else {
82            let (token_partition, _, modes) = encode_token_partition(
83                source,
84                mb_width,
85                mb_height,
86                profile,
87                segment,
88                &segment_quants,
89                &probabilities,
90                None,
91            );
92            (token_partition, probabilities, modes)
93        }
94    } else {
95        let (token_partition, _, modes) = encode_token_partition(
96            source,
97            mb_width,
98            mb_height,
99            profile,
100            segment,
101            &segment_quants,
102            &COEFFS_PROBA0,
103            None,
104        );
105        (token_partition, COEFFS_PROBA0, modes)
106    };
107    Ok(EncodedLossyCandidate {
108        base_quant: segment.quantizer[0],
109        segment: segment.clone(),
110        probabilities,
111        modes,
112        token_partition,
113    })
114}
115
116/// Finalizes the lossy candidate by choosing the best filter configuration.
117fn finalize_lossy_candidate(
118    width: usize,
119    height: usize,
120    source: &Planes,
121    mb_width: usize,
122    mb_height: usize,
123    base_quant: i32,
124    optimization_level: u8,
125    candidate: &EncodedLossyCandidate,
126) -> Result<Vec<u8>, EncoderError> {
127    let mb_count = mb_width * mb_height;
128    if !use_exhaustive_filter_search(optimization_level, mb_count) {
129        let filter = heuristic_filter(base_quant);
130        return build_candidate_vp8_frame(width, height, mb_width, mb_height, candidate, &filter);
131    }
132
133    let filters = filter_candidates(base_quant);
134    let mut best = None;
135    for filter in &filters {
136        let vp8 = build_candidate_vp8_frame(width, height, mb_width, mb_height, candidate, filter)?;
137        let distortion = yuv_sse(source, width, height, &vp8)?;
138        let replace = match &best {
139            Some((best_distortion, best_len, _)) => {
140                distortion < *best_distortion
141                    || (distortion == *best_distortion && vp8.len() < *best_len)
142            }
143            None => true,
144        };
145        if replace {
146            best = Some((distortion, vp8.len(), vp8));
147        }
148    }
149
150    best.map(|(_, _, vp8)| vp8).ok_or(EncoderError::Bitstream(
151        "lossy filter search produced no output",
152    ))
153}
154
155/// Encodes RGBA pixels to a raw lossy `VP8` frame payload with explicit options.
156pub fn encode_lossy_rgba_to_vp8_with_options(
157    width: usize,
158    height: usize,
159    rgba: &[u8],
160    options: &LossyEncodingOptions,
161) -> Result<Vec<u8>, EncoderError> {
162    validate_rgba(width, height, rgba)?;
163    validate_options(options)?;
164
165    let mb_width = (width + 15) >> 4;
166    let mb_height = (height + 15) >> 4;
167    let base_quant = base_quantizer_from_quality(options.quality);
168    let profile = lossy_search_profile(options.optimization_level);
169    let source = rgba_to_yuv420(width, height, rgba, mb_width, mb_height);
170    let candidates = build_segment_candidates(
171        &source,
172        mb_width,
173        mb_height,
174        base_quant,
175        options.optimization_level,
176    );
177    let mut best = None;
178    for segment in &candidates {
179        let candidate = encode_lossy_candidate(&source, mb_width, mb_height, &profile, segment)?;
180        let vp8 = finalize_lossy_candidate(
181            width,
182            height,
183            &source,
184            mb_width,
185            mb_height,
186            base_quant,
187            options.optimization_level,
188            &candidate,
189        )?;
190        let replace = match &best {
191            Some((best_bytes, _)) => vp8.len() < *best_bytes,
192            None => true,
193        };
194        if replace {
195            best = Some((vp8.len(), vp8));
196        }
197    }
198
199    best.map(|(_, vp8)| vp8).ok_or(EncoderError::Bitstream(
200        "lossy candidate search produced no output",
201    ))
202}
203
204/// Encodes RGBA pixels to a raw lossy `VP8` frame payload.
205pub fn encode_lossy_rgba_to_vp8(
206    width: usize,
207    height: usize,
208    rgba: &[u8],
209) -> Result<Vec<u8>, EncoderError> {
210    encode_lossy_rgba_to_vp8_with_options(width, height, rgba, &LossyEncodingOptions::default())
211}
212
213/// Encodes RGBA pixels to a still lossy WebP container with explicit options.
214pub fn encode_lossy_rgba_to_webp_with_options(
215    width: usize,
216    height: usize,
217    rgba: &[u8],
218    options: &LossyEncodingOptions,
219) -> Result<Vec<u8>, EncoderError> {
220    encode_lossy_rgba_to_webp_with_options_and_exif(width, height, rgba, options, None)
221}
222
223/// Encodes RGBA pixels to a still lossy WebP container with explicit options and EXIF.
224pub fn encode_lossy_rgba_to_webp_with_options_and_exif(
225    width: usize,
226    height: usize,
227    rgba: &[u8],
228    options: &LossyEncodingOptions,
229    exif: Option<&[u8]>,
230) -> Result<Vec<u8>, EncoderError> {
231    let vp8 = encode_lossy_rgba_to_vp8_with_options(width, height, rgba, options)?;
232    wrap_still_webp(
233        StillImageChunk {
234            fourcc: *b"VP8 ",
235            payload: &vp8,
236            width,
237            height,
238            has_alpha: false,
239        },
240        exif,
241    )
242}
243
244/// Encodes RGBA pixels to a still lossy WebP container.
245pub fn encode_lossy_rgba_to_webp(
246    width: usize,
247    height: usize,
248    rgba: &[u8],
249) -> Result<Vec<u8>, EncoderError> {
250    encode_lossy_rgba_to_webp_with_options(width, height, rgba, &LossyEncodingOptions::default())
251}
252
253/// Encodes an [`ImageBuffer`] to a still lossy WebP container with explicit options.
254pub fn encode_lossy_image_to_webp_with_options(
255    image: &ImageBuffer,
256    options: &LossyEncodingOptions,
257) -> Result<Vec<u8>, EncoderError> {
258    encode_lossy_image_to_webp_with_options_and_exif(image, options, None)
259}
260
261/// Encodes an [`ImageBuffer`] to a still lossy WebP container with explicit options and EXIF.
262pub fn encode_lossy_image_to_webp_with_options_and_exif(
263    image: &ImageBuffer,
264    options: &LossyEncodingOptions,
265    exif: Option<&[u8]>,
266) -> Result<Vec<u8>, EncoderError> {
267    encode_lossy_rgba_to_webp_with_options_and_exif(
268        image.width,
269        image.height,
270        &image.rgba,
271        options,
272        exif,
273    )
274}
275
276/// Encodes an [`ImageBuffer`] to a still lossy WebP container.
277pub fn encode_lossy_image_to_webp(image: &ImageBuffer) -> Result<Vec<u8>, EncoderError> {
278    encode_lossy_image_to_webp_with_options(image, &LossyEncodingOptions::default())
279}
280
281#[cfg(test)]
282mod tests {
283    use super::*;
284    use crate::decoder::decode_lossy_vp8_to_yuv;
285
286    fn sample_rgba() -> (usize, usize, Vec<u8>) {
287        let width = 19;
288        let height = 17;
289        let mut rgba = vec![0u8; width * height * 4];
290        for y in 0..height {
291            for x in 0..width {
292                let offset = (y * width + x) * 4;
293                rgba[offset] = (x as u8).saturating_mul(12);
294                rgba[offset + 1] = (y as u8).saturating_mul(13);
295                rgba[offset + 2] = ((x + y) as u8).saturating_mul(7);
296                rgba[offset + 3] = 0xff;
297            }
298        }
299        (width, height, rgba)
300    }
301
302    #[test]
303    fn internal_reconstruction_matches_decoder_output() {
304        let (width, height, rgba) = sample_rgba();
305        let mb_width = (width + 15) >> 4;
306        let mb_height = (height + 15) >> 4;
307        let options = LossyEncodingOptions::default();
308        let base_quant = base_quantizer_from_quality(options.quality);
309        let profile = lossy_search_profile(options.optimization_level);
310        let source = rgba_to_yuv420(width, height, &rgba, mb_width, mb_height);
311        let segment = disabled_segment_config(mb_width * mb_height, clipped_quantizer(base_quant));
312        let candidate =
313            encode_lossy_candidate(&source, mb_width, mb_height, &profile, &segment).unwrap();
314        let partition0 = encode_partition0(
315            mb_width,
316            mb_height,
317            base_quant as u8,
318            &segment,
319            &FilterConfig {
320                simple: false,
321                level: 0,
322                sharpness: 0,
323            },
324            &candidate.probabilities,
325            &candidate.modes,
326        );
327        let vp8 = build_vp8_frame(width, height, &partition0, &candidate.token_partition).unwrap();
328        let decoded = decode_lossy_vp8_to_yuv(&vp8).unwrap();
329        let (_, reconstructed, _) = encode_token_partition(
330            &source,
331            mb_width,
332            mb_height,
333            &profile,
334            &segment,
335            &build_segment_quantizers(&segment),
336            &candidate.probabilities,
337            None,
338        );
339        assert_eq!(decoded.y, reconstructed.y);
340        assert_eq!(decoded.u, reconstructed.u);
341        assert_eq!(decoded.v, reconstructed.v);
342    }
343
344    #[test]
345    fn mode_search_prefers_vertical_prediction_for_repeated_top_rows() {
346        let mb_width = 1;
347        let mb_height = 2;
348        let mut source = empty_reconstructed_planes(mb_width, mb_height);
349        let mut reconstructed = empty_reconstructed_planes(mb_width, mb_height);
350
351        for row in 0..16 {
352            for col in 0..16 {
353                let value = (col as u8).saturating_mul(9);
354                reconstructed.y[row * reconstructed.y_stride + col] = value;
355                source.y[(16 + row) * source.y_stride + col] = value;
356            }
357        }
358
359        for row in 0..8 {
360            for col in 0..8 {
361                let u = (32 + col * 7) as u8;
362                let v = (96 + col * 5) as u8;
363                reconstructed.u[row * reconstructed.uv_stride + col] = u;
364                reconstructed.v[row * reconstructed.uv_stride + col] = v;
365                source.u[(8 + row) * source.uv_stride + col] = u;
366                source.v[(8 + row) * source.uv_stride + col] = v;
367            }
368        }
369
370        let quant = build_quant_matrices(base_quantizer_from_quality(90));
371        let rd = build_rd_multipliers(&quant);
372        let profile = lossy_search_profile(MAX_LOSSY_OPTIMIZATION_LEVEL);
373        let top_modes = [B_DC_PRED; 4];
374        let left_modes = [B_DC_PRED; 4];
375        let top_context = NonZeroContext::default();
376        let left_context = NonZeroContext::default();
377        let mode = choose_macroblock_mode(
378            &source,
379            &mut reconstructed,
380            0,
381            1,
382            &profile,
383            &quant,
384            &rd,
385            &COEFFS_PROBA0,
386            &top_context,
387            &left_context,
388            &top_modes,
389            &left_modes,
390        );
391        assert!(matches!(mode.luma, V_PRED | B_PRED));
392        assert_eq!(mode.chroma, V_PRED);
393    }
394
395    #[test]
396    fn segment_candidates_include_segmented_plan_for_mixed_activity() {
397        let width = 64;
398        let height = 32;
399        let mb_width = (width + 15) >> 4;
400        let mb_height = (height + 15) >> 4;
401        let mut rgba = vec![0u8; width * height * 4];
402        for y in 0..height {
403            for x in 0..width {
404                let offset = (y * width + x) * 4;
405                let (r, g, b) = if x < width / 2 {
406                    (0x80, 0x80, 0x80)
407                } else {
408                    (
409                        ((x * 17 + y * 3) & 0xff) as u8,
410                        ((x * 5 + y * 11) & 0xff) as u8,
411                        ((x * 13 + y * 7) & 0xff) as u8,
412                    )
413                };
414                rgba[offset] = r;
415                rgba[offset + 1] = g;
416                rgba[offset + 2] = b;
417                rgba[offset + 3] = 0xff;
418            }
419        }
420
421        let source = rgba_to_yuv420(width, height, &rgba, mb_width, mb_height);
422        let candidates = build_segment_candidates(
423            &source,
424            mb_width,
425            mb_height,
426            13,
427            MAX_LOSSY_OPTIMIZATION_LEVEL,
428        );
429
430        assert!(candidates.iter().any(|candidate| candidate.use_segment));
431        assert!(candidates
432            .iter()
433            .filter(|candidate| candidate.use_segment)
434            .any(|candidate| candidate.segments.iter().any(|&segment| segment != 0)));
435    }
436
437    #[test]
438    fn segment_candidates_can_use_more_than_two_segments() {
439        let width = 96;
440        let height = 64;
441        let mb_width = (width + 15) >> 4;
442        let mb_height = (height + 15) >> 4;
443        let mut rgba = vec![0u8; width * height * 4];
444        for y in 0..height {
445            for x in 0..width {
446                let offset = (y * width + x) * 4;
447                let band = x / 24;
448                let value = match band {
449                    0 => 96,
450                    1 => ((x * 3 + y * 5) & 0xff) as u8,
451                    2 => ((x * 9 + y * 13) & 0xff) as u8,
452                    _ => ((x * 17 + y * 29) & 0xff) as u8,
453                };
454                rgba[offset] = value;
455                rgba[offset + 1] = value.wrapping_add((band * 17) as u8);
456                rgba[offset + 2] = value.wrapping_add((band * 33) as u8);
457                rgba[offset + 3] = 0xff;
458            }
459        }
460
461        let source = rgba_to_yuv420(width, height, &rgba, mb_width, mb_height);
462        let candidates = build_segment_candidates(
463            &source,
464            mb_width,
465            mb_height,
466            13,
467            MAX_LOSSY_OPTIMIZATION_LEVEL,
468        );
469
470        assert!(candidates.iter().any(|candidate| {
471            candidate.use_segment && candidate.segments.iter().copied().max().unwrap_or(0) >= 2
472        }));
473    }
474}
webp_rust/encoder/lossy/api.rs

webp_rust/encoder/lossy/
api.rs
