jxl_frame/
header.rs

1use crate::Result;
2use jxl_bitstream::{Bitstream, U};
3use jxl_image::{BitDepth, Extensions, ImageHeader, SizeHeader};
4use jxl_oxide_common::{define_bundle, Bundle, Name};
5
6define_bundle! {
7    /// Frame header.
8    #[derive(Debug)]
9    pub struct FrameHeader ctx(headers: &ImageHeader) error(crate::Error) {
10        all_default: ty(Bool) default(true),
11        pub frame_type: ty(Bundle(FrameType)) cond(!all_default) default(FrameType::RegularFrame),
12        pub encoding: ty(Bundle(Encoding)) cond(!all_default) default(Encoding::VarDct),
13        pub flags: ty(Bundle(FrameFlags)) cond(!all_default),
14        pub do_ycbcr: ty(Bool) cond(!all_default && !headers.metadata.xyb_encoded),
15        encoded_color_channels:
16            ty(u(0))
17            cond(false)
18            default({
19                let acutally_grayscale = encoding == Encoding::Modular
20                    && !do_ycbcr
21                    && !headers.metadata.xyb_encoded
22                    && headers.metadata.grayscale();
23                if acutally_grayscale { 1 } else { 3 }
24            }),
25        pub jpeg_upsampling: ty(Array[u(2)]; 3) cond(do_ycbcr && !flags.use_lf_frame()),
26        pub upsampling: ty(U32(1, 2, 4, 8)) cond(!all_default && !flags.use_lf_frame()) default(1),
27        pub ec_upsampling:
28            ty(Vec[U32(1, 2, 4, 8)]; headers.metadata.ec_info.len())
29            cond(!all_default && !flags.use_lf_frame())
30            default(vec![1; headers.metadata.ec_info.len()]),
31        pub group_size_shift: ty(u(2)) cond(encoding == Encoding::Modular) default(1),
32        pub x_qm_scale:
33            ty(u(3))
34            cond(!all_default && headers.metadata.xyb_encoded && encoding == Encoding::VarDct)
35            default(Self::compute_default_xqms(encoding, headers.metadata.xyb_encoded)),
36        pub b_qm_scale:
37            ty(u(3))
38            cond(!all_default && headers.metadata.xyb_encoded && encoding == Encoding::VarDct)
39            default(2),
40        pub passes:
41            ty(Bundle(Passes))
42            cond(!all_default && frame_type != FrameType::ReferenceOnly),
43        pub lf_level: ty(1 + u(2)) cond(frame_type == FrameType::LfFrame) default(0),
44        pub have_crop: ty(Bool) cond(!all_default && frame_type != FrameType::LfFrame) default(false),
45        pub x0:
46            ty(U32(u(8), 256 + u(11), 2304 + u(14), 18688 + u(30)); UnpackSigned)
47            cond(have_crop && frame_type != FrameType::ReferenceOnly),
48        pub y0:
49            ty(U32(u(8), 256 + u(11), 2304 + u(14), 18688 + u(30)); UnpackSigned)
50            cond(have_crop && frame_type != FrameType::ReferenceOnly),
51        pub width:
52            ty(U32(u(8), 256 + u(11), 2304 + u(14), 18688 + u(30)))
53            cond(have_crop)
54            default(headers.size.width),
55        pub height:
56            ty(U32(u(8), 256 + u(11), 2304 + u(14), 18688 + u(30)))
57            cond(have_crop)
58            default(headers.size.height),
59        pub blending_info:
60            ty(Bundle(BlendingInfo))
61            ctx((
62                !headers.metadata.ec_info.is_empty(),
63                None,
64                CanvasSizeParams {
65                    have_crop,
66                    x0,
67                    y0,
68                    width,
69                    height,
70                    size: &headers.size,
71                },
72            ))
73            cond(!all_default && frame_type.is_normal_frame()),
74        pub ec_blending_info:
75            ty(Vec[Bundle(BlendingInfo)]; headers.metadata.ec_info.len())
76            ctx((
77                !headers.metadata.ec_info.is_empty(),
78                Some(blending_info.mode),
79                CanvasSizeParams {
80                    have_crop,
81                    x0,
82                    y0,
83                    width,
84                    height,
85                    size: &headers.size,
86                },
87            ))
88            cond(!all_default && frame_type.is_normal_frame()),
89        pub duration:
90            ty(U32(0, 1, u(8), u(32)))
91            cond(!all_default && frame_type.is_normal_frame() && headers.metadata.animation.is_some())
92            default(0),
93        pub timecode:
94            ty(u(32))
95            cond(!all_default && frame_type.is_normal_frame() && headers.metadata.animation.as_ref().map(|a| a.have_timecodes).unwrap_or(false))
96            default(0),
97        pub is_last:
98            ty(Bool)
99            cond(!all_default && frame_type.is_normal_frame())
100            default(frame_type == FrameType::RegularFrame),
101        pub save_as_reference:
102            ty(u(2))
103            cond(!all_default && frame_type != FrameType::LfFrame && !is_last)
104            default(0),
105        pub resets_canvas:
106            ty(Bool)
107            cond(false)
108            default(Self::resets_canvas(
109                blending_info.mode,
110                CanvasSizeParams {
111                    have_crop,
112                    x0,
113                    y0,
114                    width,
115                    height,
116                    size: &headers.size,
117                },
118            )),
119        pub save_before_ct:
120            ty(Bool)
121            cond(
122                !all_default && (
123                    frame_type == FrameType::ReferenceOnly || (
124                        resets_canvas &&
125                        (!is_last && (duration == 0 || save_as_reference != 0) && frame_type != FrameType::LfFrame)
126                    )
127                )
128            )
129            default(!frame_type.is_normal_frame()),
130        pub name: ty(Bundle(Name)) cond(!all_default),
131        pub restoration_filter: ty(Bundle(RestorationFilter)) ctx(encoding) cond(!all_default),
132        pub extensions: ty(Bundle(Extensions)) cond(!all_default),
133        pub bit_depth: ty(Bundle(BitDepth)) cond(false) default(headers.metadata.bit_depth),
134    }
135
136    #[derive(Debug)]
137    pub struct Passes error(crate::Error) {
138        pub num_passes: ty(U32(1, 2, 3, 4 + u(3))) default(1),
139        pub num_ds: ty(U32(0, 1, 2, 3 + u(1))) cond(num_passes != 1) default(0),
140        pub shift: ty(Vec[u(2)]; num_passes - 1) cond(num_passes != 1) default(vec![0; num_passes as usize - 1]),
141        pub downsample: ty(Vec[U32(1, 2, 4, 8)]; num_ds) cond(num_passes != 1) default(vec![1; num_ds as usize]),
142        pub last_pass: ty(Vec[U32(0, 1, 2, u(3))]; num_ds) cond(num_passes != 1) default(vec![0; num_ds as usize]),
143    }
144
145    #[derive(Debug)]
146    pub struct BlendingInfo ctx(context: (bool, Option<BlendMode>, CanvasSizeParams<'_>)) error(crate::Error) {
147        pub mode: ty(Bundle(BlendMode)),
148        pub alpha_channel:
149            ty(U32(0, 1, 2, 3 + u(3)))
150            cond(context.0 && (mode == BlendMode::Blend || mode == BlendMode::MulAdd))
151            default(0),
152        pub clamp:
153            ty(Bool)
154            cond((context.0 && (mode == BlendMode::Blend || mode == BlendMode::MulAdd)) || mode == BlendMode::Mul)
155            default(false),
156        pub source:
157            ty(u(2))
158            cond(!FrameHeader::resets_canvas(context.1.unwrap_or(mode), context.2))
159            default(0),
160    }
161
162    #[derive(Debug)]
163    pub struct RestorationFilter ctx(encoding: Encoding) error(crate::Error) {
164        all_default: ty(Bool) default(true),
165        pub gab: ty(Bundle(crate::filter::Gabor)) cond(!all_default),
166        pub epf: ty(Bundle(crate::filter::EdgePreservingFilter)) cond(!all_default),
167        pub extensions: ty(Bundle(Extensions)) cond(!all_default),
168    }
169}
170
171#[derive(Copy, Clone)]
172struct CanvasSizeParams<'a> {
173    have_crop: bool,
174    x0: i32,
175    y0: i32,
176    width: u32,
177    height: u32,
178    size: &'a SizeHeader,
179}
180
181impl FrameHeader {
182    fn test_full_image(canvas_size: CanvasSizeParams) -> bool {
183        let CanvasSizeParams {
184            x0,
185            y0,
186            width,
187            height,
188            size,
189            ..
190        } = canvas_size;
191
192        if x0 > 0 || y0 > 0 {
193            return false;
194        }
195
196        let right = x0 as i64 + (width as i64);
197        let bottom = y0 as i64 + (height as i64);
198        (right >= size.width as i64) && (bottom >= size.height as i64)
199    }
200
201    fn resets_canvas(blending_mode: BlendMode, canvas_size: CanvasSizeParams) -> bool {
202        blending_mode == BlendMode::Replace
203            && (!canvas_size.have_crop || Self::test_full_image(canvas_size))
204    }
205
206    fn compute_default_xqms(encoding: Encoding, xyb_encoded: bool) -> u32 {
207        if xyb_encoded && encoding == Encoding::VarDct {
208            3
209        } else {
210            2
211        }
212    }
213
214    /// Returns whether this frame is a keyframe that should be displayed.
215    #[inline]
216    pub fn is_keyframe(&self) -> bool {
217        self.frame_type.is_normal_frame() && (self.is_last || self.duration != 0)
218    }
219
220    #[inline]
221    pub fn can_reference(&self) -> bool {
222        !self.is_last
223            && (self.duration == 0 || self.save_as_reference != 0)
224            && self.frame_type != FrameType::LfFrame
225    }
226
227    pub fn sample_width(&self, upsampling: u32) -> u32 {
228        let &Self {
229            mut width,
230            lf_level,
231            ..
232        } = self;
233
234        if upsampling > 1 {
235            width = width.div_ceil(upsampling);
236        }
237        if lf_level > 0 {
238            let div = 1u32 << (3 * lf_level);
239            width = (width + div - 1) >> (3 * lf_level);
240        }
241
242        width
243    }
244
245    pub fn sample_height(&self, upsampling: u32) -> u32 {
246        let &Self {
247            mut height,
248            lf_level,
249            ..
250        } = self;
251
252        if upsampling > 1 {
253            height = height.div_ceil(upsampling);
254        }
255        if lf_level > 0 {
256            let div = 1u32 << (3 * lf_level);
257            height = (height + div - 1) >> (3 * lf_level);
258        }
259
260        height
261    }
262
263    pub fn color_sample_width(&self) -> u32 {
264        self.sample_width(self.upsampling)
265    }
266
267    pub fn color_sample_height(&self) -> u32 {
268        self.sample_height(self.upsampling)
269    }
270
271    /// Returns the number of channels actually encoded in the frame.
272    #[inline]
273    pub fn encoded_color_channels(&self) -> usize {
274        self.encoded_color_channels as usize
275    }
276
277    pub fn num_groups(&self) -> u32 {
278        let width = self.color_sample_width();
279        let height = self.color_sample_height();
280        let group_dim = self.group_dim();
281
282        let hgroups = width.div_ceil(group_dim);
283        let vgroups = height.div_ceil(group_dim);
284
285        hgroups * vgroups
286    }
287
288    pub fn num_lf_groups(&self) -> u32 {
289        let width = self.color_sample_width();
290        let height = self.color_sample_height();
291        let lf_group_dim = self.lf_group_dim();
292
293        let hgroups = width.div_ceil(lf_group_dim);
294        let vgroups = height.div_ceil(lf_group_dim);
295
296        hgroups * vgroups
297    }
298
299    pub fn group_dim(&self) -> u32 {
300        128 << self.group_size_shift
301    }
302
303    pub fn groups_per_row(&self) -> u32 {
304        let group_dim = self.group_dim();
305        self.color_sample_width().div_ceil(group_dim)
306    }
307
308    pub fn lf_group_dim(&self) -> u32 {
309        self.group_dim() * 8
310    }
311
312    pub fn lf_groups_per_row(&self) -> u32 {
313        let lf_group_dim = self.lf_group_dim();
314        self.color_sample_width().div_ceil(lf_group_dim)
315    }
316
317    pub fn group_size_for(&self, group_idx: u32) -> (u32, u32) {
318        self.size_for(self.group_dim(), group_idx)
319    }
320
321    pub fn lf_group_size_for(&self, lf_group_idx: u32) -> (u32, u32) {
322        self.size_for(self.lf_group_dim(), lf_group_idx)
323    }
324
325    fn size_for(&self, group_dim: u32, group_idx: u32) -> (u32, u32) {
326        let width = self.color_sample_width();
327        let height = self.color_sample_height();
328        let full_rows = height / group_dim;
329        let rows_remainder = height % group_dim;
330        let full_cols = width / group_dim;
331        let cols_remainder = width % group_dim;
332
333        let stride = full_cols + (cols_remainder > 0) as u32;
334        let row = group_idx / stride;
335        let col = group_idx % stride;
336
337        let group_width = if col >= full_cols {
338            cols_remainder
339        } else {
340            group_dim
341        };
342        let group_height = if row >= full_rows {
343            rows_remainder
344        } else {
345            group_dim
346        };
347        (group_width, group_height)
348    }
349
350    pub fn lf_group_idx_from_group_idx(&self, group_idx: u32) -> u32 {
351        let groups_per_row = self.groups_per_row();
352        let lf_group_col = (group_idx % groups_per_row) / 8;
353        let lf_group_row = (group_idx / groups_per_row) / 8;
354        lf_group_col + lf_group_row * self.lf_groups_per_row()
355    }
356
357    pub fn group_idx_from_coord(&self, x: u32, y: u32) -> Option<u32> {
358        let shift = 7 + self.group_size_shift;
359        let group_x = x >> shift;
360        let group_y = y >> shift;
361        let group_idx = group_y * self.groups_per_row() + group_x;
362        if group_x >= self.groups_per_row() || group_idx >= self.num_groups() {
363            None
364        } else {
365            Some(group_idx)
366        }
367    }
368
369    pub fn is_group_collides_region(&self, group_idx: u32, region: (u32, u32, u32, u32)) -> bool {
370        let group_dim = self.group_dim();
371        let group_per_row = self.groups_per_row();
372        let group_left = (group_idx % group_per_row) * group_dim;
373        let group_top = (group_idx / group_per_row) * group_dim;
374        is_aabb_collides(region, (group_left, group_top, group_dim, group_dim))
375    }
376
377    pub fn is_lf_group_collides_region(
378        &self,
379        lf_group_idx: u32,
380        region: (u32, u32, u32, u32),
381    ) -> bool {
382        let lf_group_dim = self.lf_group_dim();
383        let lf_group_per_row = self.lf_groups_per_row();
384        let group_left = (lf_group_idx % lf_group_per_row) * lf_group_dim;
385        let group_top = (lf_group_idx / lf_group_per_row) * lf_group_dim;
386        is_aabb_collides(region, (group_left, group_top, lf_group_dim, lf_group_dim))
387    }
388}
389
390#[derive(Debug, Default, PartialEq, Eq, Copy, Clone)]
391#[repr(u8)]
392pub enum FrameType {
393    #[default]
394    RegularFrame = 0,
395    LfFrame,
396    ReferenceOnly,
397    SkipProgressive,
398}
399
400impl FrameType {
401    pub fn is_normal_frame(&self) -> bool {
402        matches!(self, Self::RegularFrame | Self::SkipProgressive)
403    }
404
405    pub fn is_progressive_frame(&self) -> bool {
406        matches!(self, Self::RegularFrame | Self::LfFrame)
407    }
408}
409
410impl<Ctx> Bundle<Ctx> for FrameType {
411    type Error = crate::Error;
412
413    fn parse(bitstream: &mut Bitstream, _ctx: Ctx) -> Result<Self> {
414        Ok(match bitstream.read_bits(2)? {
415            0 => Self::RegularFrame,
416            1 => Self::LfFrame,
417            2 => Self::ReferenceOnly,
418            3 => Self::SkipProgressive,
419            _ => unreachable!(),
420        })
421    }
422}
423
424#[derive(Debug, Default, PartialEq, Eq, Copy, Clone)]
425#[repr(u8)]
426pub enum Encoding {
427    #[default]
428    VarDct = 0,
429    Modular,
430}
431
432impl<Ctx> Bundle<Ctx> for Encoding {
433    type Error = crate::Error;
434
435    fn parse(bitstream: &mut Bitstream, _ctx: Ctx) -> Result<Self> {
436        Ok(match bitstream.read_bits(1)? {
437            0 => Self::VarDct,
438            1 => Self::Modular,
439            _ => unreachable!(),
440        })
441    }
442}
443
444#[derive(Debug, Default, PartialEq, Eq, Copy, Clone)]
445pub struct FrameFlags(u64);
446
447impl FrameFlags {
448    const NOISE: u64 = 0x1;
449    const PATCHES: u64 = 0x2;
450    const SPLINES: u64 = 0x10;
451    const USE_LF_FRAME: u64 = 0x20;
452    const SKIP_ADAPTIVE_LF_SMOOTHING: u64 = 0x80;
453
454    pub fn noise(&self) -> bool {
455        self.0 & Self::NOISE != 0
456    }
457
458    pub fn patches(&self) -> bool {
459        self.0 & Self::PATCHES != 0
460    }
461
462    pub fn splines(&self) -> bool {
463        self.0 & Self::SPLINES != 0
464    }
465
466    pub fn use_lf_frame(&self) -> bool {
467        self.0 & Self::USE_LF_FRAME != 0
468    }
469
470    pub fn skip_adaptive_lf_smoothing(&self) -> bool {
471        self.0 & Self::SKIP_ADAPTIVE_LF_SMOOTHING != 0
472    }
473}
474
475impl<Ctx> Bundle<Ctx> for FrameFlags {
476    type Error = crate::Error;
477
478    fn parse(bitstream: &mut Bitstream, _ctx: Ctx) -> Result<Self> {
479        Ok(Self(bitstream.read_u64()?))
480    }
481}
482
483#[derive(Debug, Default, PartialEq, Eq, Copy, Clone)]
484#[repr(u8)]
485pub enum BlendMode {
486    #[default]
487    Replace = 0,
488    Add = 1,
489    Blend = 2,
490    MulAdd = 3,
491    Mul = 4,
492}
493
494impl<Ctx> Bundle<Ctx> for BlendMode {
495    type Error = crate::Error;
496
497    fn parse(bitstream: &mut Bitstream, _ctx: Ctx) -> Result<Self> {
498        Ok(match bitstream.read_u32(0, 1, 2, 3 + U(2))? {
499            0 => Self::Replace,
500            1 => Self::Add,
501            2 => Self::Blend,
502            3 => Self::MulAdd,
503            4 => Self::Mul,
504            value => {
505                return Err(jxl_bitstream::Error::InvalidEnum {
506                    name: "BlendMode",
507                    value,
508                }
509                .into())
510            }
511        })
512    }
513}
514
515impl BlendMode {
516    #[inline]
517    pub fn use_alpha(self) -> bool {
518        matches!(self, Self::Blend | Self::MulAdd)
519    }
520}
521
522fn is_aabb_collides(rect0: (u32, u32, u32, u32), rect1: (u32, u32, u32, u32)) -> bool {
523    let (x0, y0, w0, h0) = rect0;
524    let (x1, y1, w1, h1) = rect1;
525    (x0 < x1 + w1) && (x0 + w0 > x1) && (y0 < y1 + h1) && (y0 + h0 > y1)
526}