leodos_protocols/application/compression/
spectral.rs1pub fn upshift(image: &mut [i32], u: u32) {
19 for sample in image.iter_mut() {
20 *sample <<= u;
21 }
22}
23
24pub fn downshift(image: &mut [i32], d: u32) {
27 if d == 0 {
28 return;
29 }
30 let half = 1i32 << (d - 1);
31 for sample in image.iter_mut() {
32 *sample = (*sample + half) >> d;
33 }
34}
35
36pub const IWT_LEVELS: usize = 5;
38
39pub fn iwt_forward_single(x: &[i32], l: &mut [i32], h: &mut [i32]) {
44 let n = x.len();
45 if n <= 1 {
46 if n == 1 {
47 l[0] = x[0];
48 }
49 return;
50 }
51
52 let p = (n + 1) / 2;
53 let q = n / 2;
54
55 for k in 0..q {
56 let is_last = k == q - 1 && n % 2 == 0;
57 h[k] = if is_last {
58 x[2 * k + 1] - x[2 * k]
59 } else {
60 x[2 * k + 1] - ((x[2 * k] + x[2 * k + 2]) / 2)
61 };
62 }
63
64 for k in 0..p {
65 l[k] = if k == 0 {
66 x[0] + ((h[0] + 1) / 2)
67 } else if k == p - 1 && n % 2 != 0 {
68 x[2 * k] + ((h[k - 1] + 1) / 2)
69 } else {
70 x[2 * k] + ((h[k - 1] + h[k] + 2) / 4)
71 };
72 }
73}
74
75pub fn iwt_inverse_single(l: &[i32], h: &[i32], x: &mut [i32]) {
77 let p = l.len();
78 let q = h.len();
79 let n = p + q;
80
81 if n <= 1 {
82 if n == 1 {
83 x[0] = l[0];
84 }
85 return;
86 }
87
88 for k in 0..p {
89 x[2 * k] = if k == 0 {
90 l[0] - ((h[0] + 1) / 2)
91 } else if k == p - 1 && n % 2 != 0 {
92 l[k] - ((h[k - 1] + 1) / 2)
93 } else {
94 l[k] - ((h[k - 1] + h[k] + 2) / 4)
95 };
96 }
97
98 for k in 0..q {
99 let is_last = k == q - 1 && n % 2 == 0;
100 x[2 * k + 1] = if is_last {
101 h[k] + x[2 * k]
102 } else {
103 h[k] + ((x[2 * k] + x[2 * k + 2]) / 2)
104 };
105 }
106}
107
108pub fn iwt_forward_5(input: &[i32], output: &mut [i32], scratch: &mut [i32]) {
113 let n = input.len();
114 output[..n].copy_from_slice(input);
115
116 let mut current_len = n;
117 for _ in 0..IWT_LEVELS {
118 if current_len <= 1 {
119 break;
120 }
121 let p = (current_len + 1) / 2;
122 let q = current_len / 2;
123
124 let (src, _) = output.split_at(current_len);
125 let src_copy: heapless::Vec<i32, 4096> = src.iter().copied().collect();
126
127 let (scratch_l, scratch_h) = scratch.split_at_mut(p);
128 iwt_forward_single(&src_copy[..current_len], scratch_l, &mut scratch_h[..q]);
129
130 output[..p].copy_from_slice(&scratch[..p]);
131 output[p..p + q].copy_from_slice(&scratch[p..p + q]);
132
133 current_len = p;
134 }
135}
136
137pub fn iwt_inverse_5(input: &[i32], output: &mut [i32], scratch: &mut [i32]) {
139 let n = input.len();
140 output[..n].copy_from_slice(input);
141
142 let mut sizes = [0usize; IWT_LEVELS + 1];
143 sizes[0] = n;
144 for i in 0..IWT_LEVELS {
145 sizes[i + 1] = (sizes[i] + 1) / 2;
146 if sizes[i] <= 1 {
147 break;
148 }
149 }
150
151 let levels_used = sizes.iter().take_while(|&&s| s > 1).count().min(IWT_LEVELS);
152
153 for level in (0..levels_used).rev() {
154 let current_len = sizes[level];
155 let p = (current_len + 1) / 2;
156 let q = current_len / 2;
157
158 scratch[..p].copy_from_slice(&output[..p]);
159 scratch[p..p + q].copy_from_slice(&output[p..p + q]);
160
161 iwt_inverse_single(&scratch[..p], &scratch[p..p + q], &mut output[..current_len]);
162 }
163}
164
165#[derive(Debug, Copy, Clone, Eq, PartialEq)]
167pub enum SpectralTransform {
168 Identity,
170 Iwt,
172}
173
174pub fn transform_spectral(
179 image: &mut [i32],
180 nx: usize,
181 ny: usize,
182 nz: usize,
183 transform: SpectralTransform,
184) {
185 match transform {
186 SpectralTransform::Identity => {}
187 SpectralTransform::Iwt => {
188 let mut col = [0i32; 4096];
189 let mut out = [0i32; 4096];
190 let mut scratch = [0i32; 4096];
191 assert!(nz <= 4096);
192
193 for y in 0..ny {
194 for x in 0..nx {
195 for z in 0..nz {
196 col[z] = image[z * nx * ny + y * nx + x];
197 }
198 iwt_forward_5(&col[..nz], &mut out[..nz], &mut scratch[..nz]);
199 for z in 0..nz {
200 image[z * nx * ny + y * nx + x] = out[z];
201 }
202 }
203 }
204 }
205 }
206}
207
208pub fn inverse_transform_spectral(
210 image: &mut [i32],
211 nx: usize,
212 ny: usize,
213 nz: usize,
214 transform: SpectralTransform,
215) {
216 match transform {
217 SpectralTransform::Identity => {}
218 SpectralTransform::Iwt => {
219 let mut col = [0i32; 4096];
220 let mut out = [0i32; 4096];
221 let mut scratch = [0i32; 4096];
222 assert!(nz <= 4096);
223
224 for y in 0..ny {
225 for x in 0..nx {
226 for z in 0..nz {
227 col[z] = image[z * nx * ny + y * nx + x];
228 }
229 iwt_inverse_5(&col[..nz], &mut out[..nz], &mut scratch[..nz]);
230 for z in 0..nz {
231 image[z * nx * ny + y * nx + x] = out[z];
232 }
233 }
234 }
235 }
236 }
237}
238
239#[cfg(test)]
240mod tests {
241 use super::*;
242
243 #[test]
244 fn upshift_downshift_roundtrip() {
245 let original = [10, -5, 127, 0, -128];
246 let mut data: [i32; 5] = original;
247 upshift(&mut data, 3);
248 for (o, d) in original.iter().zip(data.iter()) {
249 assert_eq!(*d, *o << 3);
250 }
251 downshift(&mut data, 3);
252 assert_eq!(data, original);
253 }
254
255 #[test]
256 fn downshift_rounds() {
257 let mut data = [3i32, 5, 7, -3];
258 downshift(&mut data, 1);
259 assert_eq!(data, [2, 3, 4, -1]);
260 }
261
262 #[test]
263 fn iwt_single_level_roundtrip() {
264 let input = [10, 20, 30, 40, 50, 60, 70, 80];
265 let n = input.len();
266 let p = (n + 1) / 2;
267 let q = n / 2;
268 let mut l = [0i32; 4];
269 let mut h = [0i32; 4];
270 iwt_forward_single(&input, &mut l[..p], &mut h[..q]);
271
272 let mut recovered = [0i32; 8];
273 iwt_inverse_single(&l[..p], &h[..q], &mut recovered[..n]);
274 assert_eq!(recovered, input);
275 }
276
277 #[test]
278 fn iwt_single_level_odd_length() {
279 let input = [10, 20, 30, 40, 50];
280 let n = input.len();
281 let p = (n + 1) / 2;
282 let q = n / 2;
283 let mut l = [0i32; 3];
284 let mut h = [0i32; 2];
285 iwt_forward_single(&input, &mut l[..p], &mut h[..q]);
286
287 let mut recovered = [0i32; 5];
288 iwt_inverse_single(&l[..p], &h[..q], &mut recovered[..n]);
289 assert_eq!(recovered, input);
290 }
291
292 #[test]
293 fn iwt_five_level_roundtrip() {
294 let input: [i32; 32] = core::array::from_fn(|i| (i as i32) * 7 - 100);
295 let mut output = [0i32; 32];
296 let mut scratch = [0i32; 32];
297 iwt_forward_5(&input, &mut output, &mut scratch);
298
299 let mut recovered = [0i32; 32];
300 iwt_inverse_5(&output, &mut recovered, &mut scratch);
301 assert_eq!(recovered, input);
302 }
303
304 #[test]
305 fn iwt_five_level_short_sequence() {
306 let input = [42i32, -7, 100];
307 let mut output = [0i32; 3];
308 let mut scratch = [0i32; 3];
309 iwt_forward_5(&input, &mut output, &mut scratch);
310
311 let mut recovered = [0i32; 3];
312 iwt_inverse_5(&output, &mut recovered, &mut scratch);
313 assert_eq!(recovered, input);
314 }
315
316 #[test]
317 fn spectral_transform_identity() {
318 let mut image = [1, 2, 3, 4, 5, 6, 7, 8];
319 let original = image;
320 transform_spectral(&mut image, 2, 2, 2, SpectralTransform::Identity);
321 assert_eq!(image, original);
322 }
323
324 #[test]
325 fn spectral_transform_iwt_roundtrip() {
326 let nx = 2;
327 let ny = 2;
328 let nz = 8;
329 let mut image: [i32; 32] = core::array::from_fn(|i| (i as i32) * 3 + 1);
330 let original = image;
331
332 transform_spectral(&mut image, nx, ny, nz, SpectralTransform::Iwt);
333 assert_ne!(image, original);
334
335 inverse_transform_spectral(&mut image, nx, ny, nz, SpectralTransform::Iwt);
336 assert_eq!(image, original);
337 }
338
339 #[test]
340 fn iwt_decorrelates_constant_bands() {
341 let input = [42i32; 16];
342 let mut output = [0i32; 16];
343 let mut scratch = [0i32; 16];
344 iwt_forward_5(&input, &mut output, &mut scratch);
345 assert_eq!(output[0], 42);
346 for &h in &output[1..] {
347 assert_eq!(h, 0);
348 }
349 }
350}