exr/compression/piz/
wavelet.rs

1
2//! Wavelet encoding and decoding.
3// see https://github.com/AcademySoftwareFoundation/openexr/blob/8cd1b9210855fa4f6923c1b94df8a86166be19b1/OpenEXR/IlmImf/ImfWav.cpp
4
5use crate::error::IoResult;
6use crate::math::Vec2;
7
8#[allow(unused)]
9#[inline]
10pub fn encode(buffer: &mut [u16], count: Vec2<usize>, size: Vec2<usize>, max_value: u16) -> IoResult<()> {
11    if is_14_bit(max_value) { encode_14_or_16_bit(buffer, count, size, true) }
12    else { encode_14_or_16_bit(buffer, count, size, false) }
13}
14
15#[allow(unused)]
16#[inline]
17pub fn encode_14_or_16_bit(
18    buffer: &mut [u16],
19    Vec2(count_x, count_y): Vec2<usize>,
20    Vec2(offset_x, offset_y): Vec2<usize>,
21    is_14_bit: bool // true if maximum buffer[i] value < (1 << 14)
22) -> IoResult<()>
23{
24    let count = count_x.min(count_y);
25    let encode = if is_14_bit { encode_14bit } else { encode_16bit }; // assume inlining and constant propagation
26
27    let mut p: usize = 1; // TODO i32?
28    let mut p2: usize = 2; // TODO what is p??
29
30    while p2 <= count {
31
32        let mut position_y = 0;
33        let end_y = 0 + offset_y * (count_y - p2);
34        let (offset1_x, offset1_y) = (offset_x * p, offset_y * p);
35        let (offset2_x, offset2_y) = (offset_x * p2, offset_y * p2);
36
37        // y-loop
38        while position_y <= end_y { // TODO: for py in (index..ey).nth(offset_2.0)
39
40            let mut position_x = position_y;
41            let end_x = position_x + offset_x * (count_x - p2);
42
43            // x-loop
44            while position_x <= end_x {
45                let pos_right = position_x + offset1_x;
46                let pos_top = position_x + offset1_y;
47                let pos_top_right = pos_top + offset1_x;
48
49                assert!(position_x < buffer.len());
50                assert!(pos_right < buffer.len());
51                assert!(pos_top < buffer.len());
52                assert!(pos_top_right < buffer.len());
53
54                if is_14_bit {
55                    debug_assert!(self::is_14_bit(buffer[position_x]));
56                    debug_assert!(self::is_14_bit(buffer[pos_right]));
57                }
58
59                let (center, right) = encode(buffer[position_x], buffer[pos_right]);
60                let (top, top_right) = encode(buffer[pos_top], buffer[pos_top_right]);
61
62                let (center, top) = encode(center, top);
63                let (right, top_right) = encode(right, top_right);
64
65                buffer[position_x] = center; // TODO rustify
66                buffer[pos_top] = top;
67                buffer[pos_right] = right;
68                buffer[pos_top_right] = top_right;
69
70                position_x += offset2_x;
71            }
72
73            // encode remaining odd pixel column
74            if count_x & p != 0 {
75                let pos_top = position_x + offset1_y;
76                let (center, top) = encode(buffer[position_x], buffer[pos_top]);
77
78                buffer[position_x] = center;
79                buffer[pos_top] = top;
80            }
81
82            position_y += offset2_y;
83        }
84
85        // encode possibly remaining odd row
86        if count_y & p != 0 {
87            let mut position_x = position_y;
88            let end_x = position_y + offset_x * (count_x - p2);
89
90            while position_x <= end_x {
91                let pos_right = position_x + offset1_x;
92                let (center, right) = encode(buffer[position_x], buffer[pos_right]);
93
94                buffer[pos_right] = right;
95                buffer[position_x] = center;
96
97                position_x += offset2_x;
98            }
99        }
100
101        p = p2;
102        p2 <<= 1;
103    }
104
105    Ok(())
106}
107
108#[inline]
109pub fn decode(buffer: &mut [u16], count: Vec2<usize>, size: Vec2<usize>, max_value: u16) -> IoResult<()> {
110    if is_14_bit(max_value) { decode_14_or_16_bit(buffer, count, size, true) }
111    else { decode_14_or_16_bit(buffer, count, size, false) }
112}
113
114#[inline]
115pub fn decode_14_or_16_bit(
116    buffer: &mut [u16],
117    Vec2(count_x, count_y): Vec2<usize>,
118    Vec2(offset_x, offset_y): Vec2<usize>,
119    is_14_bit: bool // true if maximum buffer[i] value < (1 << 14)
120) -> IoResult<()>
121{
122    let count = count_x.min(count_y);
123    let decode = if is_14_bit { decode_14bit } else { decode_16bit }; // assume inlining and constant propagation
124
125    let mut p: usize = 1; // TODO i32?
126    let mut p2: usize; // TODO i32?
127
128    // search max level
129    while p <= count {
130        p <<= 1;
131    }
132
133    p >>= 1;
134    p2 = p;
135    p >>= 1;
136
137    while p >= 1 {
138
139        let mut position_y = 0;
140        let end_y = 0 + offset_y * (count_y - p2);
141
142        let (offset1_x, offset1_y) = (offset_x * p, offset_y * p);
143        let (offset2_x, offset2_y) = (offset_x * p2, offset_y * p2);
144
145        debug_assert_ne!(offset_x, 0, "offset should not be zero");
146        debug_assert_ne!(offset_y, 0, "offset should not be zero");
147
148        while position_y <= end_y {
149            let mut position_x = position_y;
150            let end_x = position_x + offset_x * (count_x - p2);
151
152            while position_x <= end_x {
153                let pos_right = position_x + offset1_x;
154                let pos_top = position_x + offset1_y;
155                let pos_top_right = pos_top + offset1_x;
156
157                assert!(position_x < buffer.len());
158                assert!(pos_right < buffer.len());
159                assert!(pos_top < buffer.len());
160                assert!(pos_top_right < buffer.len());
161
162                let (center, top) = decode(buffer[position_x], buffer[pos_top]);
163                let (right, top_right) = decode(buffer[pos_right], buffer[pos_top_right]);
164
165                let (center, right) = decode(center, right);
166                let (top, top_right) = decode(top, top_right);
167
168                buffer[position_x] = center; // TODO rustify
169                buffer[pos_top] = top;
170                buffer[pos_right] = right;
171                buffer[pos_top_right] = top_right;
172
173                position_x += offset2_x;
174            }
175
176            // decode last odd remaining x value
177            if count_x & p != 0 {
178                let pos_top = position_x + offset1_y;
179                let (center, top) = decode(buffer[position_x], buffer[pos_top]);
180
181                buffer[position_x] = center;
182                buffer[pos_top] = top;
183            }
184
185            position_y += offset2_y;
186        }
187
188        // decode remaining odd row
189        if count_y & p != 0 {
190            let mut position_x = position_y;
191            let end_x = position_x + offset_x * (count_x - p2);
192
193            while position_x <= end_x {
194                let pos_right = position_x + offset1_x;
195                let (center, right) = decode(buffer[position_x], buffer[pos_right]);
196
197                buffer[position_x] = center;
198                buffer[pos_right] = right;
199
200                position_x += offset2_x;
201            }
202        }
203
204        p2 = p;
205        p >>= 1;
206    }
207
208    Ok(())
209}
210
211#[inline]
212fn is_14_bit(value: u16) -> bool {
213    value < (1 << 14)
214}
215
216/// Untransformed data values should be less than (1 << 14).
217#[inline]
218#[allow(unused)]
219fn encode_14bit(a: u16, b: u16) -> (u16, u16) {
220    let (a, b) = (a as i16, b as i16);
221
222    let m = (a + b) >> 1;
223    let d = a - b;
224
225    (m as u16, d as u16) // TODO explicitly wrap?
226}
227
228#[inline]
229#[allow(unused)]
230fn decode_14bit(l: u16, h: u16) -> (u16, u16) {
231    let (l, h) = (l as i16, h as i16);
232
233    let hi = h as i32;
234    let ai = l as i32 + (hi & 1) + (hi >> 1);
235
236    let a = ai as i16; // TODO explicitly wrap?
237    let b = (ai - hi) as i16; // TODO explicitly wrap?
238
239    (a as u16, b as u16) // TODO explicitly wrap?
240}
241
242
243const BIT_COUNT: i32 = 16;
244const OFFSET: i32 = 1 << (BIT_COUNT - 1);
245const MOD_MASK: i32 = (1 << BIT_COUNT) - 1;
246
247#[inline]
248fn encode_16bit(a: u16, b: u16) -> (u16, u16) {
249    let (a, b) = (a as i32, b as i32);
250
251    let a_offset = (a + OFFSET) & MOD_MASK;
252    let mut m = (a_offset + b) >> 1;
253    let d = a_offset - b;
254
255    if d < 0 { m = (m + OFFSET) & MOD_MASK; }
256    let d = d & MOD_MASK;
257
258    (m as u16, d as u16) // TODO explicitly wrap?
259}
260
261#[inline]
262fn decode_16bit(l: u16, h: u16) -> (u16, u16) {
263    let (m, d) = (l as i32, h as i32);
264
265    let b = (m - (d >> 1)) & MOD_MASK;
266    let a = (d + b - OFFSET) & MOD_MASK;
267
268    (a as u16, b as u16) // TODO explicitly wrap?
269}
270
271
272
273#[cfg(test)]
274mod test {
275    use crate::math::Vec2;
276    use crate::compression::piz::wavelet::is_14_bit;
277
278    #[test]
279    fn roundtrip_14_bit_values(){
280        let data = [
281            (13, 54), (3, 123), (423, 53), (1, 23), (23, 515), (513, 43),
282            (16374, 16381), (16284, 3), (2, 1), (0, 0), (0, 4), (3, 0)
283        ];
284
285        for &values in &data {
286            let (l, h) = super::encode_14bit(values.0, values.1);
287            let result = super::decode_14bit(l, h);
288            assert_eq!(values, result);
289        }
290    }
291
292    #[test]
293    fn roundtrip_16_bit_values(){
294        let data = [
295            (13, 54), (3, 123), (423, 53), (1, 23), (23, 515), (513, 43),
296            (16385, 56384), (18384, 36384), (2, 1), (0, 0), (0, 4), (3, 0)
297        ];
298
299        for &values in &data {
300            let (l, h) = super::encode_16bit(values.0, values.1);
301            let result = super::decode_16bit(l, h);
302            assert_eq!(values, result);
303        }
304    }
305
306    #[test]
307    fn roundtrip_14bit_image(){
308        let data: [u16; 6 * 4] = [
309            13, 54, 3, 123, 423, 53,
310            1, 23, 23, 515, 513, 43,
311            16374, 16381, 16284, 3, 2, 1,
312            0, 0, 0, 4, 3, 0,
313        ];
314
315        let max = *data.iter().max().unwrap();
316        debug_assert!(is_14_bit(max));
317
318        let mut transformed = data.clone();
319
320        super::encode(&mut transformed, Vec2(6, 4), Vec2(1,6), max).unwrap();
321        super::decode(&mut transformed, Vec2(6, 4), Vec2(1,6), max).unwrap();
322
323        assert_eq!(data, transformed);
324    }
325
326    #[test]
327    fn roundtrip_16bit_image(){
328        let data: [u16; 6 * 4] = [
329            13, 54, 3, 123, 423, 53,
330            1, 23, 23, 515, 513, 43,
331            16385, 56384, 18384, 36384, 2, 1,
332            0, 0, 0, 4, 3, 0,
333        ];
334
335        let max = *data.iter().max().unwrap();
336        debug_assert!(!is_14_bit(max));
337
338        let mut transformed = data.clone();
339
340        super::encode(&mut transformed, Vec2(6, 4), Vec2(1,6), max).unwrap();
341        super::decode(&mut transformed, Vec2(6, 4), Vec2(1,6), max).unwrap();
342
343        assert_eq!(data, transformed);
344    }
345
346    /// inspired by https://github.com/AcademySoftwareFoundation/openexr/blob/master/OpenEXR/IlmImfTest/testWav.cpp
347    #[test]
348    fn ground_truth(){
349        test_size(1, 1);
350        test_size(2, 2);
351        test_size(32, 32);
352        test_size(1024, 16);
353        test_size(16, 1024);
354        test_size(997, 37);
355        test_size(37, 997);
356        test_size(1024, 1024);
357        test_size(997, 997);
358
359        fn test_size(x: usize, y: usize) {
360            let xy = Vec2(x, y);
361            roundtrip(noise_14bit(xy), xy);
362            roundtrip(noise_16bit(xy), xy);
363            roundtrip(solid(xy, 0), xy);
364            roundtrip(solid(xy, 1), xy);
365            roundtrip(solid(xy, 0xffff), xy);
366            roundtrip(solid(xy, 0x3fff), xy);
367            roundtrip(solid(xy, 0x3ffe), xy);
368            roundtrip(solid(xy, 0x3fff), xy);
369            roundtrip(solid(xy, 0xfffe), xy);
370            roundtrip(solid(xy, 0xffff), xy);
371            roundtrip(verticals(xy, 0xffff), xy);
372            roundtrip(verticals(xy, 0x3fff), xy);
373            roundtrip(horizontals(xy, 0xffff), xy);
374            roundtrip(horizontals(xy, 0x3fff), xy);
375            roundtrip(diagonals(xy, 0xffff), xy);
376            roundtrip(diagonals(xy, 0x3fff), xy);
377        }
378
379        fn roundtrip(data: Vec<u16>, size: Vec2<usize>){
380            assert_eq!(data.len(), size.area());
381
382            let max = *data.iter().max().unwrap();
383            let offset = Vec2(1, size.0);
384
385            let mut transformed = data.clone();
386            super::encode(&mut transformed, size, offset, max).unwrap();
387            super::decode(&mut transformed, size, offset, max).unwrap();
388
389            assert_eq!(data, transformed);
390        }
391
392        fn noise_14bit(size: Vec2<usize>) -> Vec<u16> {
393            (0..size.area()).map(|_| (rand::random::<i32>() & 0x3fff) as u16).collect()
394        }
395
396        fn noise_16bit(size: Vec2<usize>) -> Vec<u16> {
397            (0..size.area()).map(|_| rand::random::<u16>()).collect()
398        }
399
400        fn solid(size: Vec2<usize>, value: u16) -> Vec<u16> {
401            vec![value; size.area()]
402        }
403
404        fn verticals(size: Vec2<usize>, max_value: u16) -> Vec<u16> {
405            std::iter::repeat_with(|| (0 .. size.0).map(|x| if x & 1 != 0 { 0 } else { max_value }))
406                .take(size.1).flatten().collect()
407        }
408
409        fn horizontals(size: Vec2<usize>, max_value: u16) -> Vec<u16> {
410            (0 .. size.1)
411                .flat_map(|y| std::iter::repeat(if y & 1 != 0 { 0 } else { max_value }).take(size.0))
412                .collect()
413        }
414
415        fn diagonals(size: Vec2<usize>, max_value: u16) -> Vec<u16> {
416            (0 .. size.1).flat_map(|y| {
417                (0 .. size.0).map(move |x| if (x + y) & 1 != 0 { 0 } else { max_value })
418            }).collect()
419        }
420
421    }
422}