jxl-render 0.12.4

JPEG XL image renderer, part of jxl-oxide
Documentation 
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
131
132

use jxl_grid::{AlignedGrid, AllocTracker, SharedSubgrid};

use crate::Region;
use crate::util::PaddedGrid;

pub fn upsample(
    grid: SharedSubgrid<f32>,
    out_region: &mut Region,
    image_header: &jxl_image::ImageHeader,
    factor: u32,
    tracker: Option<&AllocTracker>,
) -> crate::Result<Option<AlignedGrid<f32>>> {
    let metadata = &image_header.metadata;

    let mut out = None;
    let mut grid = grid;

    let up8 = factor / 3;
    let last_up = factor % 3;
    for _ in 0..up8 {
        out = Some(upsample_inner::<8, 210>(
            grid,
            &metadata.up8_weight,
            tracker,
        )?);
        grid = out.as_ref().unwrap().as_subgrid();
    }
    out = match last_up {
        1 => Some(upsample_inner::<2, 15>(
            grid,
            &metadata.up2_weight,
            tracker,
        )?),
        2 => Some(upsample_inner::<4, 55>(
            grid,
            &metadata.up4_weight,
            tracker,
        )?),
        _ => out,
    };
    *out_region = out_region.upsample(factor);
    Ok(out)
}

fn upsample_inner<const K: usize, const NW: usize>(
    grid: SharedSubgrid<f32>,
    weights: &[f32; NW],
    tracker: Option<&AllocTracker>,
) -> crate::Result<AlignedGrid<f32>> {
    assert!((K == 2 && NW == 15) || (K == 4 && NW == 55) || (K == 8 && NW == 210));
    let grid_width = grid.width();
    let grid_height = grid.height();
    let frame_width = grid_width << K.ilog2();
    let frame_height = grid_height << K.ilog2();

    // 5x5 kernel
    const PADDING: usize = 2;
    let mut padded = PaddedGrid::with_alloc_tracker(grid_width, grid_height, PADDING, tracker)?;
    let padded_width = grid_width + PADDING * 2;

    let padded_buf = padded.buf_padded_mut();
    for y in 0..grid.height() {
        let row = grid.get_row(y);
        padded_buf[(y + PADDING) * padded_width + PADDING..][..grid_width].copy_from_slice(row);
    }
    padded.mirror_edges_padding();

    // For K = 8:
    // 0  1  2  3  3h 2h 1h 0h
    // 4  5  6  7  7h 6h 5h 4h
    // 8  9  a  b  bh ah 9h 8h
    // c  d  e  f  fh eh dh ch
    // cv dv ev fv f' e' d' c'
    // 8v 9v av bv b' a' 9' 8'
    // 4v 5v 6v 7v 7' 6' 5' 4'
    // 0v 1v 2v 3v 3' 2' 1' 0'
    let mut weights_quarter = vec![[0.0f32; 25]; K * K / 4];
    let mut weight_idx = 0usize;
    let mat_n = K / 2;
    for y in 0..5 * mat_n {
        let mat_y = y / 5;
        let ky = y % 5;
        for x in y..5 * mat_n {
            let mat_x = x / 5;
            let kx = x % 5;
            let w = weights[weight_idx];
            weight_idx += 1;

            weights_quarter[mat_y * mat_n + mat_x][ky * 5 + kx] = w;
            weights_quarter[mat_x * mat_n + mat_y][kx * 5 + ky] = w;
        }
    }

    let mut grid = AlignedGrid::with_alloc_tracker(frame_width, frame_height, tracker)?;
    let padded_buf = padded.buf_padded();
    let grid_buf = grid.buf_mut();
    for y in 0..frame_height {
        let ref_y = y / K;
        let mat_y = (y % K).min(K - y % K - 1);
        let flip_v = y % K >= mat_n;
        for x in 0..frame_width {
            let ref_x = x / K;
            let mat_x = (x % K).min(K - x % K - 1);
            let flip_h = x % K >= mat_n;

            let kernel = &weights_quarter[mat_y * mat_n + mat_x];
            let mut sum = 0.0f32;
            let mut min = f32::INFINITY;
            let mut max = -f32::INFINITY;
            for iy in 0..5 {
                let ky = if flip_v { 4 - iy } else { iy };
                for ix in 0..5 {
                    let kx = if flip_h { 4 - ix } else { ix };
                    let sample = padded_buf[(ref_y + iy) * padded_width + (ref_x + ix)];
                    sum += kernel[ky * 5 + kx] * sample;
                    min = min.min(sample);
                    max = max.max(sample);
                }
            }

            grid_buf[y * frame_width + x] = if !min.is_finite() {
                // This happens if all samples are NaN. f32::clamp would panic
                // in this case.
                f32::NAN
            } else {
                sum.clamp(min, max)
            };
        }
    }

    Ok(grid)
}