Skip to main content

docling_pdf/render/
jpx.rs

1//! `JPXDecode` images (ISO 32000-1, 7.4.9): JPEG 2000 codestreams and JP2
2//! files, decoded with `hayro-jpeg2000` (#598).
3//!
4//! A photograph stored as JPEG 2000 used to draw as a mid-gray block, and
5//! the layout model then saw a blank rectangle where docling(-parse) shows it
6//! the picture: on the reporter's NASA scans the picture box lost its
7//! confidence (0.64 instead of 0.99), ran over the caption line below, and
8//! two stacked photos separated by a caption fused into one — the caption
9//! then nested inside the picture and vanished from the Markdown. Decoded,
10//! the same pages give the same regions as the docling-parse renderer.
11//!
12//! What the filter dictionary decides and what the codestream decides
13//! (7.4.9): the colour space comes from the codestream (gray, RGB — sYCC
14//! already converted —, CMYK, or by component count for an ICC profile) when
15//! the image dictionary has no `/ColorSpace`; a dictionary `/ColorSpace`
16//! wins, and an `Indexed` one means the samples are palette indices, so the
17//! decoder is told not to resolve the codestream's own palette.
18//! `/SMaskInData` 1 or 2 takes the codestream's alpha channel as the soft
19//! mask (2 = premultiplied, which the renderer treats like 1: the blit
20//! premultiplies again, a shade darker on the edges of such an image);
21//! without it the alpha channel is dropped. `/BitsPerComponent` is the
22//! codestream's; the decoder normalizes every depth to 8 bits.
23//!
24//! A reduced decode (`target`, the renderer's `codec_reduction_shift` hint)
25//! lets the decoder stop at a lower resolution level when the image is drawn
26//! far smaller than it is; the returned size is whatever it decoded at.
27
28use hayro_jpeg2000::{ColorSpace as JpxColorSpace, DecodeSettings, DecoderContext, Image};
29
30/// A decoded JPX image: 8-bit interleaved colour samples, the alpha channel
31/// (when the codestream has one) split off.
32pub struct Decoded {
33    pub width: usize,
34    pub height: usize,
35    /// Colour components per pixel (the alpha channel excluded).
36    pub ncomp: usize,
37    /// `width * height * ncomp` bytes.
38    pub data: Vec<u8>,
39    /// The codestream's opacity channel, `width * height` bytes.
40    pub alpha: Option<Vec<u8>>,
41    /// What the codestream says its colour components are.
42    pub color: Color,
43}
44
45/// The codestream's colour interpretation, as far as the image dictionary
46/// needs it when it names no `/ColorSpace` of its own.
47#[derive(Clone, Copy, Debug, PartialEq, Eq)]
48pub enum Color {
49    Gray,
50    Rgb,
51    Cmyk,
52    /// An ICC profile or an unknown space: by component count.
53    Other,
54}
55
56/// Decode `data` (a JP2 file or a raw J2K codestream). `indexed`: the PDF
57/// colour space is `Indexed`, so a palette in the codestream must stay
58/// unresolved (the samples are the indices). `target`: a `(width, height)`
59/// the decoder may stop at (a lower resolution level), `None` for full size.
60pub fn decode(data: &[u8], indexed: bool, target: Option<(u32, u32)>) -> Result<Decoded, String> {
61    let settings = DecodeSettings {
62        resolve_palette_indices: !indexed,
63        target_resolution: target,
64        ..DecodeSettings::default()
65    };
66    let image = Image::new(data, &settings).map_err(|e| format!("jpx: {e:?}"))?;
67    let color = match image.color_space() {
68        JpxColorSpace::Gray => Color::Gray,
69        JpxColorSpace::RGB => Color::Rgb,
70        JpxColorSpace::CMYK => Color::Cmyk,
71        _ => Color::Other,
72    };
73    let mut ctx = DecoderContext::default();
74    let decoded = image.decode(&mut ctx).map_err(|e| format!("jpx: {e:?}"))?;
75    // The image's size is the decoded one: with a target resolution the
76    // header's extent is already divided by the skipped levels.
77    let (width, height) = (image.width() as usize, image.height() as usize);
78    let channels = decoded.components().len();
79    if width == 0 || height == 0 || channels == 0 {
80        return Err("jpx: empty image".into());
81    }
82    let interleaved = decoded.data_u8();
83    if interleaved.len() != width * height * channels {
84        return Err(format!(
85            "jpx: {} samples for {width}x{height}x{channels}",
86            interleaved.len()
87        ));
88    }
89    // The alpha channel, when declared, is the last one.
90    let has_alpha = image.has_alpha() && channels >= 2;
91    let ncomp = if has_alpha { channels - 1 } else { channels };
92    let (data, alpha) = if has_alpha {
93        let mut data = Vec::with_capacity(width * height * ncomp);
94        let mut alpha = Vec::with_capacity(width * height);
95        for px in interleaved.chunks_exact(channels) {
96            data.extend_from_slice(&px[..ncomp]);
97            alpha.push(px[ncomp]);
98        }
99        (data, Some(alpha))
100    } else {
101        (interleaved, None)
102    };
103    Ok(Decoded {
104        width,
105        height,
106        ncomp,
107        data,
108        alpha,
109        color,
110    })
111}
112
113#[cfg(test)]
114mod tests {
115    use super::*;
116
117    fn fixture(name: &str) -> Vec<u8> {
118        std::fs::read(
119            std::path::Path::new(env!("CARGO_MANIFEST_DIR"))
120                .join("tests/data/jpx")
121                .join(name),
122        )
123        .expect("jpx fixture")
124    }
125
126    /// The fixtures are OpenJPEG (Pillow) encodes of known gradients,
127    /// lossless (5/3 wavelet), so the samples come back exactly.
128    #[test]
129    fn gray_jp2_and_raw_codestream_decode_exactly() {
130        for name in ["gray_12x9.jp2", "gray_12x9.j2k"] {
131            let d = decode(&fixture(name), false, None).unwrap();
132            assert_eq!((d.width, d.height, d.ncomp), (12, 9, 1), "{name}");
133            assert_eq!(d.color, Color::Gray, "{name}");
134            assert!(d.alpha.is_none());
135            for y in 0..9 {
136                for x in 0..12 {
137                    assert_eq!(
138                        d.data[y * 12 + x],
139                        ((x * 21 + y * 3) % 256) as u8,
140                        "{name} ({x},{y})"
141                    );
142                }
143            }
144        }
145    }
146
147    #[test]
148    fn rgb_jp2_decodes_exactly() {
149        let d = decode(&fixture("rgb_8x6.jp2"), false, None).unwrap();
150        assert_eq!((d.width, d.height, d.ncomp), (8, 6, 3));
151        assert_eq!(d.color, Color::Rgb);
152        for y in 0..6 {
153            for x in 0..8 {
154                let px = &d.data[(y * 8 + x) * 3..][..3];
155                assert_eq!(
156                    px,
157                    [
158                        ((x * 32) % 256) as u8,
159                        ((y * 40) % 256) as u8,
160                        (((x + y) * 17) % 256) as u8
161                    ],
162                    "({x},{y})"
163                );
164            }
165        }
166    }
167
168    /// A JP2 with an opacity channel: colour and alpha come apart, alpha
169    /// row-major at the image size.
170    #[test]
171    fn alpha_channel_is_split_off() {
172        let d = decode(&fixture("rgba_8x6.jp2"), false, None).unwrap();
173        assert_eq!((d.width, d.height, d.ncomp), (8, 6, 3));
174        let alpha = d.alpha.expect("alpha channel");
175        assert_eq!(alpha.len(), 48);
176        for y in 0..6 {
177            for x in 0..8 {
178                assert_eq!(alpha[y * 8 + x], ((x * 36) % 256) as u8, "({x},{y})");
179                assert_eq!(d.data[(y * 8 + x) * 3 + 2], 128, "({x},{y})");
180            }
181        }
182    }
183
184    /// A target resolution lets the decoder stop a level early; the samples
185    /// match the size it reports.
186    #[test]
187    fn reduced_decode_reports_its_own_size() {
188        let full = decode(&fixture("gray_12x9.jp2"), false, None).unwrap();
189        let small = decode(&fixture("gray_12x9.jp2"), false, Some((6, 4))).unwrap();
190        assert_eq!(small.data.len(), small.width * small.height * small.ncomp);
191        assert!(small.width <= full.width && small.height <= full.height);
192        assert!(
193            small.width >= 3 && small.height >= 2,
194            "{}x{}",
195            small.width,
196            small.height
197        );
198    }
199
200    #[test]
201    fn garbage_is_an_error_not_a_panic() {
202        assert!(decode(b"not a jpx", false, None).is_err());
203        assert!(decode(&[], false, None).is_err());
204        let mut truncated = fixture("gray_12x9.jp2");
205        truncated.truncate(60);
206        let _ = decode(&truncated, false, None); // either way, no panic
207    }
208}