jxl_frame/data/
lf_group.rs

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
use jxl_bitstream::Bitstream;
use jxl_grid::AllocTracker;
use jxl_modular::{image::TransformedModularSubimage, MaConfig, Sample};
use jxl_oxide_common::Bundle;
use jxl_vardct::{HfMetadata, HfMetadataParams, LfCoeff, LfCoeffParams, Quantizer};

use crate::{
    filter::{EdgePreservingFilter, EpfParams},
    header::Encoding,
    FrameHeader, Result,
};

#[derive(Debug)]
pub struct LfGroupParams<'a, 'dest, 'tracker, S: Sample> {
    pub frame_header: &'a FrameHeader,
    pub quantizer: Option<&'a Quantizer>,
    pub global_ma_config: Option<&'a MaConfig>,
    pub mlf_group: Option<TransformedModularSubimage<'dest, S>>,
    pub lf_group_idx: u32,
    pub allow_partial: bool,
    pub tracker: Option<&'tracker AllocTracker>,
    pub pool: &'a jxl_threadpool::JxlThreadPool,
}

#[derive(Debug)]
pub struct LfGroup<S: Sample> {
    pub lf_coeff: Option<LfCoeff<S>>,
    pub hf_meta: Option<HfMetadata>,
    pub partial: bool,
}

impl<S: Sample> Bundle<LfGroupParams<'_, '_, '_, S>> for LfGroup<S> {
    type Error = crate::Error;

    fn parse(bitstream: &mut Bitstream, params: LfGroupParams<S>) -> Result<Self> {
        let LfGroupParams {
            frame_header,
            global_ma_config,
            mlf_group,
            lf_group_idx,
            allow_partial,
            tracker,
            pool,
            ..
        } = params;
        let (lf_width, lf_height) = frame_header.lf_group_size_for(lf_group_idx);

        let lf_coeff = (frame_header.encoding == Encoding::VarDct
            && !frame_header.flags.use_lf_frame())
        .then(|| {
            let lf_coeff_params = LfCoeffParams {
                lf_group_idx,
                lf_width,
                lf_height,
                jpeg_upsampling: frame_header.jpeg_upsampling,
                bits_per_sample: frame_header.bit_depth.bits_per_sample(),
                global_ma_config,
                allow_partial,
                tracker,
                pool,
            };
            LfCoeff::parse(bitstream, lf_coeff_params)
        })
        .transpose()?;

        if let Some(lf_coeff_inner) = &lf_coeff {
            if lf_coeff_inner.partial {
                return Ok(Self {
                    lf_coeff,
                    hf_meta: None,
                    partial: true,
                });
            }
        }

        let mut is_mlf_complete = true;
        if let Some(image) = mlf_group {
            if !image.is_empty() {
                let mut subimage = image.recursive(bitstream, global_ma_config, tracker)?;
                let mut subimage = subimage.prepare_subimage()?;
                subimage.decode(
                    bitstream,
                    1 + frame_header.num_lf_groups() + lf_group_idx,
                    allow_partial,
                )?;
                is_mlf_complete = subimage.finish(pool);
            }
        }

        let hf_meta = (frame_header.encoding == Encoding::VarDct && is_mlf_complete)
            .then(|| {
                let hf_meta_params = HfMetadataParams {
                    num_lf_groups: frame_header.num_lf_groups(),
                    lf_group_idx,
                    lf_width,
                    lf_height,
                    jpeg_upsampling: frame_header.jpeg_upsampling,
                    bits_per_sample: frame_header.bit_depth.bits_per_sample(),
                    global_ma_config,
                    epf: match &frame_header.restoration_filter.epf {
                        EdgePreservingFilter::Disabled => None,
                        EdgePreservingFilter::Enabled(EpfParams {
                            sharp_lut, sigma, ..
                        }) => Some((sigma.quant_mul, *sharp_lut)),
                    },
                    quantizer_global_scale: params.quantizer.unwrap().global_scale,
                    tracker,
                    pool,
                };
                HfMetadata::parse(bitstream, hf_meta_params)
            })
            .transpose();
        match hf_meta {
            Err(e) if e.unexpected_eof() && allow_partial => {
                tracing::debug!("Decoded partial HfMeta");
                Ok(Self {
                    lf_coeff,
                    hf_meta: None,
                    partial: true,
                })
            }
            Err(e) => Err(e.into()),
            Ok(hf_meta) => Ok(Self {
                lf_coeff,
                hf_meta,
                partial: !is_mlf_complete,
            }),
        }
    }
}