otf_pixels_codec_jxl/
encoder.rs

1//! The JPEG XL encoder: lossy VarDCT by default, lossless Modular on
2//! request (and, until VarDCT carries it, for images with alpha).
3//!
4//! Lossless:
5//!
6//! The image becomes integer channels (grey or YCoCg colour, then alpha),
7//! each predicted by the clamped gradient and coded per channel with prefix
8//! codes fitted to its residuals. An image of one 256x256 group is a single
9//! section; a larger one is split into groups, each its own section, behind
10//! a table of contents, all sharing the codes of the global section.
11
12use otf_pixels_core::{
13    ChannelLayout, EncodeOptions, Encoder, ImageDescriptor, PixelsError, Result, SampleKind, Sink,
14};
15
16use crate::bits::BitWriter;
17use crate::entropy::{EntropyCode, HybridUint, Token, with_runs};
18use crate::headers::{ImageInfo, write_image_header, write_modular_frame_header, write_toc};
19use crate::modular::{Channel, Tree, forward_ycocg, residuals, write_header};
20use crate::predict::{Predictor, Region};
21
22/// libjxl's cutoffs for the weighted predictor's error (`kWPFixedDC`).
23const WP_CUTOFFS: [i32; 33] = [
24    -500, -392, -255, -191, -127, -95, -63, -47, -31, -23, -15, -11, -7, -4, -3, -1, 0, 1, 3, 5, 7,
25    11, 15, 23, 31, 47, 63, 95, 127, 191, 255, 392, 500,
26];
27
28/// The side of a group, the unit a large image is split into.
29const GROUP: usize = 256;
30/// The side of an LF group, which a Modular frame without squeeze leaves
31/// empty.
32const LF_GROUP: usize = GROUP * 8;
33/// The largest dimension the size header can declare.
34const MAX_DIMENSION: u32 = 1 << 30;
35
36/// Encodes a JPEG XL codestream.
37#[derive(Debug)]
38pub struct JxlEncoder {
39    state: Option<State>,
40    /// `None` for lossless; otherwise the Butteraugli-style distance.
41    distance: Option<f32>,
42    /// The EXIF block for an `Exif` box, if any.
43    exif: Option<Vec<u8>>,
44    /// The XMP packet for an `xml ` box, if any.
45    xmp: Option<Vec<u8>>,
46}
47
48impl Default for JxlEncoder {
49    fn default() -> Self {
50        Self::new()
51    }
52}
53
54#[derive(Debug)]
55struct State {
56    descriptor: ImageDescriptor,
57    pixels: Vec<u8>,
58    rows_written: u32,
59}
60
61impl JxlEncoder {
62    /// A lossy encoder at distance 1, visually lossless.
63    #[must_use]
64    pub const fn new() -> Self {
65        Self {
66            state: None,
67            distance: Some(1.0),
68            exif: None,
69            xmp: None,
70        }
71    }
72
73    /// A lossless encoder.
74    #[must_use]
75    pub const fn lossless() -> Self {
76        Self {
77            state: None,
78            distance: None,
79            exif: None,
80            xmp: None,
81        }
82    }
83
84    /// A lossy encoder at `distance`: 1 is visually lossless, 0.5 finer,
85    /// up to 25 for the coarsest.
86    #[must_use]
87    pub const fn with_distance(distance: f32) -> Self {
88        Self {
89            state: None,
90            distance: Some(distance),
91            exif: None,
92            xmp: None,
93        }
94    }
95
96    /// An encoder configured from generic encode options: lossless if
97    /// `lossless`, otherwise lossy at the distance libjxl's `cjxl` gives
98    /// `quality` (90 is distance 1).
99    #[must_use]
100    pub fn from_options(options: &EncodeOptions) -> Self {
101        if options.lossless {
102            return Self::lossless();
103        }
104        Self::with_distance(distance_for_quality(options.quality))
105    }
106}
107
108/// libjxl's quality-to-distance mapping.
109fn distance_for_quality(quality: u8) -> f32 {
110    let q = f32::from(quality.clamp(1, 100));
111    if q >= 30.0 {
112        0.1 + (100.0 - q) * 0.09
113    } else {
114        53.0 / 3000.0 * q * q - 23.0 / 20.0 * q + 25.0
115    }
116}
117
118impl JxlEncoder {}
119
120/// The header's view of `format`.
121fn info_of(descriptor: &ImageDescriptor) -> Result<ImageInfo> {
122    let format = descriptor.pixel;
123    let bit_depth = match format.sample_kind() {
124        SampleKind::U8 => 8,
125        SampleKind::U16 => 16,
126        _ => {
127            return Err(PixelsError::unsupported(format!(
128                "JPEG XL encoding takes 8- or 16-bit samples; got {format}. Convert first."
129            )));
130        }
131    };
132    Ok(ImageInfo {
133        width: descriptor.width,
134        height: descriptor.height,
135        bit_depth,
136        grey: matches!(
137            format.layout(),
138            ChannelLayout::Gray | ChannelLayout::GrayAlpha
139        ),
140        alpha: format.layout().has_alpha(),
141        xyb: false,
142    })
143}
144
145/// The image's channels: grey or R, G, B, then alpha.
146fn channels_of(descriptor: &ImageDescriptor, pixels: &[u8]) -> Vec<Channel> {
147    let format = descriptor.pixel;
148    let (width, height) = (descriptor.width as usize, descriptor.height as usize);
149    let count = format.channels();
150    let wide = format.sample_kind() == SampleKind::U16;
151    let mut channels = vec![Channel::new(width, height); count];
152    let bytes = if wide { 2 } else { 1 };
153    for (i, pixel) in pixels.chunks_exact(count * bytes).enumerate() {
154        for (c, channel) in channels.iter_mut().enumerate() {
155            let at = c * bytes;
156            let value = if wide {
157                i32::from(u16::from_ne_bytes([
158                    pixel.get(at).copied().unwrap_or(0),
159                    pixel.get(at + 1).copied().unwrap_or(0),
160                ]))
161            } else {
162                i32::from(pixel.get(at).copied().unwrap_or(0))
163            };
164            if let Some(slot) = channel.samples.get_mut(i) {
165                *slot = value;
166            }
167        }
168    }
169    channels
170}
171
172/// The whole lossy codestream for `pixels`, which have no alpha.
173fn encode_lossy(descriptor: &ImageDescriptor, pixels: &[u8], distance: f32) -> Result<Vec<u8>> {
174    let mut info = info_of(descriptor)?;
175    let channels = channels_of(descriptor, pixels);
176    let max = if info.bit_depth == 16 {
177        65_535.0
178    } else {
179        255.0
180    };
181    let n = descriptor.width as usize * descriptor.height as usize;
182    let rgb: Vec<[f32; 3]> = (0..n)
183        .map(|i| {
184            let at = |c: usize| {
185                channels
186                    .get(c)
187                    .and_then(|ch| ch.samples.get(i))
188                    .copied()
189                    .unwrap_or(0) as f32
190                    / max
191            };
192            if info.grey {
193                [at(0); 3]
194            } else {
195                [at(0), at(1), at(2)]
196            }
197        })
198        .collect();
199    // XYB carries colour; the header's sRGB says what to decode it to, and
200    // eight bits is what the samples were.
201    info.grey = false;
202    info.xyb = true;
203    info.bit_depth = 8;
204    Ok(crate::vardct::encode(&rgb, &info, distance))
205}
206
207/// The whole lossless codestream for `pixels`.
208fn encode(descriptor: &ImageDescriptor, pixels: &[u8]) -> Result<Vec<u8>> {
209    let info = info_of(descriptor)?;
210    let mut channels = channels_of(descriptor, pixels);
211    let rct = !info.grey;
212    if rct {
213        forward_ycocg(&mut channels);
214    }
215    let (width, height) = (info.width as usize, info.height as usize);
216    // libjxl's `-e3` tree: the weighted predictor everywhere, contexts by
217    // its error estimate (property 15).
218    let pixels_total = width * height * channels.len();
219    let tree = Tree::fixed(
220        15,
221        &WP_CUTOFFS,
222        Predictor::Weighted,
223        pixels_total,
224        info.bit_depth,
225    );
226    let tree_tokens = tree.tokens();
227    let tree_code = EntropyCode::build(6, &tree_tokens, HybridUint::BUCKETS);
228
229    // Every group's residuals, in raster order of groups; a single group is
230    // the whole image, coded in the global section.
231    let groups_x = width.div_ceil(GROUP);
232    let groups_y = height.div_ceil(GROUP);
233    let single = groups_x * groups_y == 1;
234    let mut group_tokens: Vec<Vec<Token>> = Vec::with_capacity(groups_x * groups_y);
235    for gy in 0..groups_y {
236        for gx in 0..groups_x {
237            let (x0, y0) = (gx * GROUP, gy * GROUP);
238            let region = (x0, y0, GROUP.min(width - x0), GROUP.min(height - y0));
239            let mut tokens = Vec::with_capacity(region.2 * region.3 * channels.len());
240            for (c, channel) in channels.iter().enumerate() {
241                let region = Region {
242                    channel,
243                    x0: region.0,
244                    y0: region.1,
245                    width: region.2,
246                    height: region.3,
247                };
248                residuals(&region, c as u32, &tree, &mut tokens);
249            }
250            // Runs of exact predictions (residual 0) become copies.
251            group_tokens.push(with_runs(&tokens, 0));
252        }
253    }
254    let all: Vec<Token> = group_tokens.iter().flatten().copied().collect();
255    let code = EntropyCode::build(tree.leaves(), &all, HybridUint::DEFAULT);
256
257    // The global section: dequantization defaults, the tree and codes, the
258    // image's header and, for a single group, its residuals.
259    let mut global = BitWriter::new();
260    global.bool(true); // default LF dequantization
261    global.bool(true); // a global tree follows
262    tree_code.write_header(&mut global);
263    for &token in &tree_tokens {
264        tree_code.write(&mut global, token);
265    }
266    code.write_header(&mut global);
267    write_header(&mut global, rct);
268    let mut sections = Vec::new();
269    if single {
270        for &token in &all {
271            code.write(&mut global, token);
272        }
273        sections.push(global.finish());
274    } else {
275        sections.push(global.finish());
276        let lf_groups = width.div_ceil(LF_GROUP) * height.div_ceil(LF_GROUP);
277        // LF groups and the HF global section hold nothing in a lossless
278        // Modular frame.
279        sections.extend(std::iter::repeat_n(Vec::new(), lf_groups + 1));
280        for tokens in &group_tokens {
281            let mut group = BitWriter::new();
282            write_header(&mut group, false);
283            for &token in tokens {
284                code.write(&mut group, token);
285            }
286            sections.push(group.finish());
287        }
288    }
289
290    let mut out = BitWriter::new();
291    write_image_header(&mut out, &info);
292    write_modular_frame_header(&mut out, &info);
293    let sizes: Vec<usize> = sections.iter().map(Vec::len).collect();
294    write_toc(&mut out, &sizes);
295    let mut bytes = out.finish();
296    for section in &sections {
297        bytes.extend_from_slice(section);
298    }
299    Ok(bytes)
300}
301
302/// An ISO-BMFF box.
303fn bx(kind: &[u8; 4], payload: &[u8]) -> Vec<u8> {
304    let mut out = u32::try_from(payload.len() + 8)
305        .unwrap_or(u32::MAX)
306        .to_be_bytes()
307        .to_vec();
308    out.extend_from_slice(kind);
309    out.extend_from_slice(payload);
310    out
311}
312
313/// The codestream in the container (ISO/IEC 18181-2), beside an `Exif`
314/// box (a zero offset to the TIFF header, then the block) and an `xml ` box.
315fn contained(codestream: &[u8], exif: Option<&[u8]>, xmp: Option<&[u8]>) -> Vec<u8> {
316    let mut out = crate::decode::container::CONTAINER_SIGNATURE.to_vec();
317    out.extend(bx(b"ftyp", b"jxl \0\0\0\0jxl "));
318    if let Some(exif) = exif {
319        out.extend(bx(b"Exif", &[&[0_u8; 4][..], exif].concat()));
320    }
321    if let Some(xmp) = xmp {
322        out.extend(bx(b"xml ", xmp));
323    }
324    out.extend(bx(b"jxlc", codestream));
325    out
326}
327
328impl Encoder for JxlEncoder {
329    /// Written in an `Exif` box, which puts the file in the container.
330    fn set_exif(&mut self, exif: Option<&[u8]>) -> Result<()> {
331        if self.state.is_some() {
332            return Err(PixelsError::invalid_argument(
333                "exif",
334                "the EXIF block must be set before write_header",
335            ));
336        }
337        self.exif = exif.map(<[u8]>::to_vec);
338        Ok(())
339    }
340
341    /// Written in an `xml ` box, which puts the file in the container.
342    fn set_xmp(&mut self, xmp: Option<&[u8]>) -> Result<()> {
343        if self.state.is_some() {
344            return Err(PixelsError::invalid_argument(
345                "xmp",
346                "the XMP packet must be set before write_header",
347            ));
348        }
349        self.xmp = xmp.map(<[u8]>::to_vec);
350        Ok(())
351    }
352
353    fn write_header(&mut self, desc: &ImageDescriptor, _sink: &mut dyn Sink) -> Result<()> {
354        if self.state.is_some() {
355            return Err(PixelsError::invalid_argument(
356                "descriptor",
357                "write_header called more than once",
358            ));
359        }
360        info_of(desc)?;
361        if desc.width == 0
362            || desc.height == 0
363            || desc.width > MAX_DIMENSION
364            || desc.height > MAX_DIMENSION
365        {
366            return Err(PixelsError::unsupported(format!(
367                "JPEG XL dimensions must be 1 to {MAX_DIMENSION}; {}x{} is not",
368                desc.width, desc.height
369            )));
370        }
371        let capacity = desc
372            .byte_len()
373            .ok_or_else(|| PixelsError::malformed("jxl", "image size overflows"))?;
374        self.state = Some(State {
375            descriptor: *desc,
376            pixels: Vec::with_capacity(capacity),
377            rows_written: 0,
378        });
379        Ok(())
380    }
381
382    fn write_row(&mut self, row: &[u8], _sink: &mut dyn Sink) -> Result<()> {
383        let Some(state) = self.state.as_mut() else {
384            return Err(PixelsError::invalid_argument(
385                "row",
386                "write_row called before write_header",
387            ));
388        };
389        let expected = state.descriptor.row_bytes();
390        if row.len() != expected {
391            return Err(PixelsError::invalid_argument(
392                "row",
393                format!("row is {} bytes, expected {expected}", row.len()),
394            ));
395        }
396        if state.rows_written >= state.descriptor.height {
397            return Err(PixelsError::invalid_argument(
398                "row",
399                format!("more than {} rows written", state.descriptor.height),
400            ));
401        }
402        state.pixels.extend_from_slice(row);
403        state.rows_written += 1;
404        Ok(())
405    }
406
407    fn finish(&mut self, sink: &mut dyn Sink) -> Result<()> {
408        let Some(state) = self.state.as_ref() else {
409            return Err(PixelsError::invalid_argument(
410                "sink",
411                "finish called before write_header",
412            ));
413        };
414        if state.rows_written < state.descriptor.height {
415            return Err(PixelsError::malformed(
416                "jxl",
417                format!(
418                    "{} of {} rows were written",
419                    state.rows_written, state.descriptor.height
420                ),
421            ));
422        }
423        let bytes = match self.distance {
424            Some(distance) if !state.descriptor.pixel.layout().has_alpha() => {
425                encode_lossy(&state.descriptor, &state.pixels, distance)?
426            }
427            _ => encode(&state.descriptor, &state.pixels)?,
428        };
429        let bytes = if self.exif.is_some() || self.xmp.is_some() {
430            contained(&bytes, self.exif.as_deref(), self.xmp.as_deref())
431        } else {
432            bytes
433        };
434        sink.write_all(&bytes)?;
435        sink.flush()
436    }
437}
otf_pixels_codec_jxl/encoder.rs

otf_pixels_codec_jxl/
encoder.rs
