a_sixel/
focal.rs

1//! Use weighted pixels based on the image's spectral properties to provided
2//! weighted input to k-means.
3//!
4//! Weights are computed using the following steps:
5//! 1. Compute the saliency maps for each channel (L, a, b) using a combination
6//!    of the following:
7//!    - Spectral residual is the difference between the blurred natural log of
8//!      the amplitude calculated by the FFT and the natural log of the
9//!      amplitude reconstructed from the inverse FFT using the original phase
10//!      spectrum.
11//!    - Modulated spectral residual is the same as the above, but the amplitude
12//!      and phase have flattened mid values.
13//!    - Phase spectrum is the phase of the FFT.
14//!    - Amplitude spectrum is the amplitude of the FFT modulated to have a
15//!      flattened midrange.
16//! 2. Compute the average distance of each pixel to the centroid of the image
17//!    transformed to the range \[0, 1].
18//! 3. Compute the average local distance of each pixel to its neighbors within
19//!    a window of size ln(pixels.len()) / 2. This is used to help identify
20//!    edges and features in the image.
21//! 4. Compute the final weight for each pixel as a combination of the saliency
22//!    maps and the local distance: lerp(a, sr.max(mod_sr), p, ld) where sr is
23//!    the spectral residual, mod_sr is the modulated spectral residual, p is
24//!    the phase spectrum, a is the amplitude spectrum, and ld is the local
25//!    distance.
26
27use std::{
28    f32::consts::PI,
29    sync::atomic::Ordering,
30};
31
32use atomic_float::AtomicF32;
33use image::{
34    Rgb,
35    RgbImage,
36};
37use libblur::{
38    stack_blur_f32,
39    AnisotropicRadius,
40    BlurImageMut,
41    FastBlurChannels,
42    ThreadingPolicy,
43};
44use palette::{
45    color_difference::EuclideanDistance,
46    IntoColor,
47    Lab,
48    Srgb,
49};
50use rayon::{
51    iter::{
52        IndexedParallelIterator,
53        IntoParallelIterator,
54        IntoParallelRefIterator,
55        IntoParallelRefMutIterator,
56        ParallelIterator,
57    },
58    slice::{
59        ParallelSlice,
60        ParallelSliceMut,
61    },
62};
63use rustfft::{
64    num_complex::Complex,
65    num_traits::Zero,
66    FftPlanner,
67};
68
69use crate::{
70    bitmerge::agglomerative_merge,
71    dither::Sierra,
72    kmeans::parallel_kmeans,
73    private,
74    rgb_to_lab,
75    BitPaletteBuilder,
76    PaletteBuilder,
77    SixelEncoder,
78};
79
80pub type FocalSixelEncoderMono<D = Sierra> = SixelEncoder<FocalPaletteBuilder<2>, D>;
81pub type FocalSixelEncoder4<D = Sierra> = SixelEncoder<FocalPaletteBuilder<4>, D>;
82pub type FocalSixelEncoder8<D = Sierra> = SixelEncoder<FocalPaletteBuilder<8>, D>;
83pub type FocalSixelEncoder16<D = Sierra> = SixelEncoder<FocalPaletteBuilder<16>, D>;
84pub type FocalSixelEncoder32<D = Sierra> = SixelEncoder<FocalPaletteBuilder<32>, D>;
85pub type FocalSixelEncoder64<D = Sierra> = SixelEncoder<FocalPaletteBuilder<64>, D>;
86pub type FocalSixelEncoder128<D = Sierra> = SixelEncoder<FocalPaletteBuilder<128>, D>;
87pub type FocalSixelEncoder256<D = Sierra> = SixelEncoder<FocalPaletteBuilder<256>, D>;
88
89pub struct FocalPaletteBuilder<const PALETTE_SIZE: usize = 256>;
90
91impl<const PALETTE_SIZE: usize> private::Sealed for FocalPaletteBuilder<PALETTE_SIZE> {}
92
93impl<const PALETTE_SIZE: usize> PaletteBuilder for FocalPaletteBuilder<PALETTE_SIZE> {
94    const NAME: &'static str = "Focal";
95    const PALETTE_SIZE: usize = PALETTE_SIZE;
96
97    #[inline(never)]
98    fn build_palette(image: &RgbImage) -> Vec<Lab> {
99        let image_width = image.width();
100        let image_height = image.height();
101        let image = image.to_vec();
102
103        let mut pixels =
104            vec![<Lab>::new(0.0, 0.0, 0.0); image_width as usize * image_height as usize];
105
106        image.par_chunks(3).zip(&mut pixels).for_each(|(p, dest)| {
107            *dest = rgb_to_lab(Rgb([p[0], p[1], p[2]]));
108        });
109
110        let window_radius = (pixels.len() as f32).ln().clamp(2.0, 16.0) as u32 / 2;
111
112        let mut planner = FftPlanner::new();
113
114        let mut l_values = vec![0.0f32; pixels.len()];
115        let mut a_values = vec![0.0f32; pixels.len()];
116        let mut b_values = vec![0.0f32; pixels.len()];
117        pixels
118            .par_iter()
119            .copied()
120            .zip(&mut l_values)
121            .zip(&mut a_values)
122            .zip(&mut b_values)
123            .for_each(|(((lab, dest_l), dest_a), dest_b)| {
124                *dest_l = lab.l;
125                *dest_a = lab.a;
126                *dest_b = lab.b;
127            });
128
129        let l_saliency = compute_saliency(
130            &mut planner,
131            &l_values,
132            image_width,
133            image_height,
134            window_radius >> 1,
135        );
136
137        let a_saliency = compute_saliency(
138            &mut planner,
139            &a_values,
140            image_width,
141            image_height,
142            window_radius >> 1,
143        );
144
145        let b_saliency = compute_saliency(
146            &mut planner,
147            &b_values,
148            image_width,
149            image_height,
150            window_radius >> 1,
151        );
152
153        #[cfg(feature = "dump_l_saliency")]
154        {
155            let mut quant_l_saliency = vec![0; pixels.len()];
156
157            l_saliency
158                .spectral_residual
159                .par_iter()
160                .copied()
161                .zip(&mut quant_l_saliency)
162                .for_each(|(d, dest)| {
163                    *dest = (d * u8::MAX as f32).round() as u8;
164                });
165            dump_intermediate("l_sr_aliency", &quant_l_saliency, image_width, image_height);
166
167            l_saliency
168                .mod_spectral_residual
169                .par_iter()
170                .copied()
171                .zip(&mut quant_l_saliency)
172                .for_each(|(d, dest)| {
173                    *dest = (d * u8::MAX as f32).round() as u8;
174                });
175            dump_intermediate(
176                "l_mod_sr_saliency",
177                &quant_l_saliency,
178                image_width,
179                image_height,
180            );
181
182            l_saliency
183                .phase_spectrum
184                .par_iter()
185                .copied()
186                .zip(&mut quant_l_saliency)
187                .for_each(|(d, dest)| {
188                    *dest = (d * u8::MAX as f32).round() as u8;
189                });
190            dump_intermediate(
191                "l_phase_saliency",
192                &quant_l_saliency,
193                image_width,
194                image_height,
195            );
196
197            l_saliency
198                .amplitude_spectrum
199                .par_iter()
200                .copied()
201                .zip(&mut quant_l_saliency)
202                .for_each(|(d, dest)| {
203                    *dest = (d * u8::MAX as f32).round() as u8;
204                });
205            dump_intermediate(
206                "l_amplitude_saliency",
207                &quant_l_saliency,
208                image_width,
209                image_height,
210            );
211        }
212
213        #[cfg(feature = "dump_a_saliency")]
214        {
215            let mut quant_a_saliency = vec![0; pixels.len()];
216
217            a_saliency
218                .spectral_residual
219                .par_iter()
220                .copied()
221                .zip(&mut quant_a_saliency)
222                .for_each(|(d, dest)| {
223                    *dest = (d * u8::MAX as f32).round() as u8;
224                });
225            dump_intermediate(
226                "a_sr_saliency",
227                &quant_a_saliency,
228                image_width,
229                image_height,
230            );
231
232            a_saliency
233                .mod_spectral_residual
234                .par_iter()
235                .copied()
236                .zip(&mut quant_a_saliency)
237                .for_each(|(d, dest)| {
238                    *dest = (d * u8::MAX as f32).round() as u8;
239                });
240            dump_intermediate(
241                "a_mod_sr_saliency",
242                &quant_a_saliency,
243                image_width,
244                image_height,
245            );
246
247            a_saliency
248                .phase_spectrum
249                .par_iter()
250                .copied()
251                .zip(&mut quant_a_saliency)
252                .for_each(|(d, dest)| {
253                    *dest = (d * u8::MAX as f32).round() as u8;
254                });
255            dump_intermediate(
256                "a_phase_saliency",
257                &quant_a_saliency,
258                image_width,
259                image_height,
260            );
261
262            a_saliency
263                .amplitude_spectrum
264                .par_iter()
265                .copied()
266                .zip(&mut quant_a_saliency)
267                .for_each(|(d, dest)| {
268                    *dest = (d * u8::MAX as f32).round() as u8;
269                });
270            dump_intermediate(
271                "a_amplitude_saliency",
272                &quant_a_saliency,
273                image_width,
274                image_height,
275            );
276        }
277
278        #[cfg(feature = "dump_b_saliency")]
279        {
280            let mut quant_b_saliency = vec![0; pixels.len()];
281
282            b_saliency
283                .spectral_residual
284                .par_iter()
285                .copied()
286                .zip(&mut quant_b_saliency)
287                .for_each(|(d, dest)| {
288                    *dest = (d * u8::MAX as f32).round() as u8;
289                });
290            dump_intermediate(
291                "b_sr_saliency",
292                &quant_b_saliency,
293                image_width,
294                image_height,
295            );
296
297            b_saliency
298                .mod_spectral_residual
299                .par_iter()
300                .copied()
301                .zip(&mut quant_b_saliency)
302                .for_each(|(d, dest)| {
303                    *dest = (d * u8::MAX as f32).round() as u8;
304                });
305            dump_intermediate(
306                "b_mod_sr_saliency",
307                &quant_b_saliency,
308                image_width,
309                image_height,
310            );
311
312            b_saliency
313                .phase_spectrum
314                .par_iter()
315                .copied()
316                .zip(&mut quant_b_saliency)
317                .for_each(|(d, dest)| {
318                    *dest = (d * u8::MAX as f32).round() as u8;
319                });
320            dump_intermediate(
321                "b_phase_saliency",
322                &quant_b_saliency,
323                image_width,
324                image_height,
325            );
326
327            b_saliency
328                .amplitude_spectrum
329                .par_iter()
330                .copied()
331                .zip(&mut quant_b_saliency)
332                .for_each(|(d, dest)| {
333                    *dest = (d * u8::MAX as f32).round() as u8;
334                });
335            dump_intermediate(
336                "b_amplitude_saliency",
337                &quant_b_saliency,
338                image_width,
339                image_height,
340            );
341        }
342
343        let max_dist = ((<Lab>::max_l() - <Lab>::min_l()).powi(2)
344            + (<Lab>::max_a() - <Lab>::min_a()).powi(2)
345            + (<Lab>::max_b() - <Lab>::min_b()).powi(2))
346        .sqrt();
347
348        let mut local_dists = vec![0.0f32; pixels.len()];
349
350        (0..pixels.len())
351            .into_par_iter()
352            .zip(&mut local_dists)
353            .for_each(|(idx, dest)| {
354                let x = idx as isize % image_width as isize;
355                let y = idx as isize / image_height as isize;
356
357                let mut sum = 0.0;
358                let mut count = 0;
359                for dx in -(window_radius as isize)..=window_radius as isize {
360                    for dy in -(window_radius as isize)..=window_radius as isize {
361                        let x = x + dx;
362                        let y = y + dy;
363
364                        if (dx == 0 && dy == 0)
365                            || x < 0
366                            || y < 0
367                            || x >= image_width as isize
368                            || y >= image_height as isize
369                        {
370                            continue;
371                        }
372                        let jdx = y * image_width as isize + x;
373                        sum += pixels[jdx as usize].distance(pixels[idx]);
374                        count += 1;
375                    }
376                }
377
378                *dest = (sum / count as f32) / max_dist
379            });
380
381        #[cfg(feature = "dump_local_saliency")]
382        {
383            let mut quant_local_dists = vec![0; pixels.len()];
384            local_dists
385                .par_iter()
386                .copied()
387                .zip(&mut quant_local_dists)
388                .for_each(|(d, dest)| {
389                    *dest = (d * u8::MAX as f32).round() as u8;
390                });
391
392            dump_intermediate("local_dists", &quant_local_dists, image_width, image_height);
393        }
394
395        let mut candidates = vec![(<Lab>::new(0.0, 0.0, 0.0), 0.0); pixels.len()];
396        pixels
397            .into_par_iter()
398            .zip(l_saliency.spectral_residual)
399            .zip(l_saliency.mod_spectral_residual)
400            .zip(l_saliency.phase_spectrum)
401            .zip(l_saliency.amplitude_spectrum)
402            .zip(a_saliency.spectral_residual)
403            .zip(a_saliency.mod_spectral_residual)
404            .zip(a_saliency.phase_spectrum)
405            .zip(a_saliency.amplitude_spectrum)
406            .zip(b_saliency.spectral_residual)
407            .zip(b_saliency.mod_spectral_residual)
408            .zip(b_saliency.phase_spectrum)
409            .zip(b_saliency.amplitude_spectrum)
410            .zip(local_dists)
411            .zip(&mut candidates)
412            .for_each(
413                |(
414                    (
415                        (
416                            (
417                                (
418                                    (
419                                        (
420                                            (
421                                                ((((((lab, l_sr), l_msr), l_p), l_a), a_sr), a_msr),
422                                                a_p,
423                                            ),
424                                            a_a,
425                                        ),
426                                        b_sr,
427                                    ),
428                                    b_msr,
429                                ),
430                                b_p,
431                            ),
432                            b_a,
433                        ),
434                        local,
435                    ),
436                    dest,
437                )| {
438                    let outlier_outlier = l_msr.max(a_msr).max(b_msr);
439                    let outliers = l_sr.max(a_sr).max(b_sr).max(outlier_outlier);
440                    let edges = l_p.max(a_p).max(b_p);
441                    let blobs = l_a.max(a_a).max(b_a);
442
443                    // Lerp between blobs <=> outliers <=> edges using local distance. Local
444                    // distance is *high* for edges and features and *low* for blobs/planes of
445                    // colors. If it's neither, then it's probably important if it's an outlier.
446                    let w = if local <= 0.5 {
447                        blobs + 2.0 * local * (outliers - blobs)
448                    } else {
449                        outliers + 2.0 * (local - 0.5) * (edges - outliers)
450                    };
451
452                    *dest = (lab, w);
453                },
454            );
455
456        #[cfg(feature = "dump_weights")]
457        {
458            let mut quant_candidates = vec![0; candidates.len()];
459            candidates
460                .par_iter()
461                .copied()
462                .zip(&mut quant_candidates)
463                .for_each(|((_, w), dest)| {
464                    *dest = (w * u8::MAX as f32).round() as u8;
465                });
466
467            dump_intermediate("candidates", &quant_candidates, image_width, image_height);
468        }
469
470        let buckets = Vec::from_iter(
471            std::iter::repeat_with(|| {
472                (
473                    (
474                        AtomicF32::new(0.0),
475                        AtomicF32::new(0.0),
476                        AtomicF32::new(0.0),
477                    ),
478                    AtomicF32::new(0.0),
479                )
480            })
481            .take(1 << 21),
482        );
483
484        candidates.par_iter().for_each(|(c, w)| {
485            let rgb: Srgb = (*c).into_color();
486            let index = BitPaletteBuilder::<{ 1 << 21 }>::index(rgb.into_format());
487            let bucket = &buckets[index];
488            bucket.0 .0.fetch_add(c.l * *w, Ordering::Relaxed);
489            bucket.0 .1.fetch_add(c.a * *w, Ordering::Relaxed);
490            bucket.0 .2.fetch_add(c.b * *w, Ordering::Relaxed);
491            bucket.1.fetch_add(*w, Ordering::Relaxed);
492        });
493
494        let candidates = buckets
495            .into_par_iter()
496            .filter_map(|bucket| {
497                let count = bucket.1.load(Ordering::Relaxed);
498                if count > 0.0 {
499                    let lab = <Lab>::new(
500                        bucket.0 .0.load(Ordering::Relaxed) / count,
501                        bucket.0 .1.load(Ordering::Relaxed) / count,
502                        bucket.0 .2.load(Ordering::Relaxed) / count,
503                    );
504                    Some((lab, count))
505                } else {
506                    None
507                }
508            })
509            .collect::<Vec<_>>();
510
511        let (mut stage2_colors, mut stage2_counts) = parallel_kmeans::<512>(&candidates);
512
513        agglomerative_merge::<512, PALETTE_SIZE>(&mut stage2_colors, &mut stage2_counts)
514    }
515}
516
517struct Saliency {
518    spectral_residual: Vec<f32>,
519    mod_spectral_residual: Vec<f32>,
520    phase_spectrum: Vec<f32>,
521    amplitude_spectrum: Vec<f32>,
522}
523
524fn compute_saliency(
525    planner: &mut FftPlanner<f32>,
526    channel_values: &[f32],
527    image_width: u32,
528    image_height: u32,
529    window_radius: u32,
530) -> Saliency {
531    let buffer = {
532        let fft_row = planner.plan_fft_forward(image_width as usize);
533        let fft_col = planner.plan_fft_forward(image_height as usize);
534
535        let mut buffer = vec![Complex::zero(); channel_values.len()];
536        channel_values
537            .par_iter()
538            .copied()
539            .zip(&mut buffer)
540            .for_each(|(v, dest)| {
541                *dest = Complex { re: v, im: 0.0 };
542            });
543
544        buffer
545            .par_chunks_mut(image_width as usize)
546            .for_each(|chunk| {
547                fft_row.process(chunk);
548            });
549
550        let mut transpose = vec![Complex::zero(); buffer.len()];
551        (0..image_width as usize)
552            .into_par_iter()
553            .zip(transpose.par_chunks_mut(image_height as usize))
554            .for_each(|(x, col)| {
555                (0..image_height as usize)
556                    .into_par_iter()
557                    .zip(col)
558                    .for_each(|(y, dest)| {
559                        *dest = buffer[y * image_width as usize + x];
560                    });
561            });
562
563        transpose
564            .par_chunks_mut(image_height as usize)
565            .for_each(|chunk| {
566                fft_col.process(chunk);
567            });
568
569        (0..image_height as usize)
570            .into_par_iter()
571            .zip(buffer.par_chunks_mut(image_width as usize))
572            .for_each(|(y, row)| {
573                (0..image_width as usize)
574                    .into_par_iter()
575                    .zip(row)
576                    .for_each(|(x, dest)| {
577                        *dest = transpose[x * image_height as usize + y];
578                    });
579            });
580
581        transpose
582    };
583
584    let mut amplitude = vec![0.0f32; buffer.len()];
585    buffer
586        .par_iter()
587        .copied()
588        .zip(&mut amplitude)
589        .for_each(|(c, dest)| *dest = c.norm());
590
591    let pft = {
592        let mut ifft_buffer = vec![Complex::zero(); buffer.len()];
593
594        amplitude
595            .par_iter()
596            .copied()
597            .zip(&buffer)
598            .zip(&mut ifft_buffer)
599            .for_each(|((a, c), dest)| {
600                *dest = Complex {
601                    re: if a > f32::EPSILON { c.re / a } else { 0.0 },
602                    im: if a > f32::EPSILON { c.im / a } else { 0.0 },
603                };
604            });
605
606        let ifft_row = planner.plan_fft_inverse(image_width as usize);
607        let ifft_col = planner.plan_fft_inverse(image_height as usize);
608
609        ifft_buffer
610            .par_chunks_mut(image_height as usize)
611            .for_each(|chunk| {
612                ifft_col.process(chunk);
613            });
614
615        let mut transpose = vec![Complex::zero(); ifft_buffer.len()];
616        (0..image_width as usize)
617            .into_par_iter()
618            .zip(transpose.par_chunks_mut(image_height as usize))
619            .for_each(|(x, col)| {
620                (0..image_height as usize)
621                    .into_par_iter()
622                    .zip(col)
623                    .for_each(|(y, dest)| {
624                        *dest = ifft_buffer[y * image_width as usize + x];
625                    });
626            });
627
628        transpose
629            .par_chunks_mut(image_height as usize)
630            .for_each(|chunk| {
631                ifft_row.process(chunk);
632            });
633
634        (0..image_height as usize)
635            .into_par_iter()
636            .zip(ifft_buffer.par_chunks_mut(image_width as usize))
637            .for_each(|(y, row)| {
638                (0..image_width as usize)
639                    .into_par_iter()
640                    .zip(row)
641                    .for_each(|(x, dest)| {
642                        *dest = transpose[x * image_height as usize + y];
643                    });
644            });
645
646        let mut pft = vec![0.0f32; transpose.len()];
647
648        transpose
649            .par_iter()
650            .copied()
651            .zip(&mut pft)
652            .for_each(|(c, dest)| *dest = c.norm_sqr());
653
654        {
655            let mut pft =
656                BlurImageMut::borrow(&mut pft, image_width, image_height, FastBlurChannels::Plane);
657            stack_blur_f32(
658                &mut pft,
659                AnisotropicRadius::new(window_radius),
660                ThreadingPolicy::Adaptive,
661            )
662            .unwrap();
663        }
664
665        let max_pft = pft.par_iter().copied().reduce(|| f32::MIN, f32::max);
666        let min_pft = pft.par_iter().copied().reduce(|| f32::MAX, f32::min);
667        pft.par_iter_mut().for_each(|a| {
668            *a = (*a - min_pft) / (max_pft - min_pft).max(f32::EPSILON);
669        });
670
671        pft
672    };
673
674    let mut phase = vec![0.0f32; buffer.len()];
675    buffer
676        .into_par_iter()
677        .zip(&mut phase)
678        .for_each(|(c, dest)| *dest = c.arg());
679
680    let mut mod_amplitude = vec![0.0f32; phase.len()];
681    let max_amplitude = amplitude.par_iter().copied().reduce(|| f32::MIN, f32::max);
682    let min_amplitude = amplitude.par_iter().copied().reduce(|| f32::MAX, f32::min);
683    amplitude
684        .par_iter()
685        .copied()
686        .zip(&mut mod_amplitude)
687        .for_each(|(a, dest)| {
688            let x = ((a - min_amplitude) / (max_amplitude - min_amplitude).max(f32::EPSILON) - 0.5)
689                * 2.0;
690            *dest = ((x * PI / 2.0).sin().powi(15) / 2.0 + 0.5) * (max_amplitude - min_amplitude)
691                + min_amplitude
692        });
693
694    let asr = {
695        let mut ifft_buffer = vec![Complex::zero(); phase.len()];
696        mod_amplitude
697            .par_iter()
698            .copied()
699            .zip(&phase)
700            .zip(&mut ifft_buffer)
701            .for_each(|((a, p), dest)| {
702                *dest = Complex {
703                    re: a * p.cos(),
704                    im: a * p.sin(),
705                };
706            });
707
708        let ifft_row = planner.plan_fft_inverse(image_width as usize);
709        let ifft_col = planner.plan_fft_inverse(image_height as usize);
710        ifft_buffer
711            .par_chunks_mut(image_height as usize)
712            .for_each(|chunk| {
713                ifft_col.process(chunk);
714            });
715
716        let mut transpose = vec![Complex::zero(); ifft_buffer.len()];
717        (0..image_width as usize)
718            .into_par_iter()
719            .zip(transpose.par_chunks_mut(image_height as usize))
720            .for_each(|(x, col)| {
721                (0..image_height as usize)
722                    .into_par_iter()
723                    .zip(col)
724                    .for_each(|(y, dest)| {
725                        *dest = ifft_buffer[y * image_width as usize + x];
726                    });
727            });
728
729        transpose
730            .par_chunks_mut(image_height as usize)
731            .for_each(|chunk| {
732                ifft_row.process(chunk);
733            });
734
735        (0..image_height as usize)
736            .into_par_iter()
737            .zip(ifft_buffer.par_chunks_mut(image_width as usize))
738            .for_each(|(y, row)| {
739                (0..image_width as usize)
740                    .into_par_iter()
741                    .zip(row)
742                    .for_each(|(x, dest)| {
743                        *dest = transpose[x * image_height as usize + y];
744                    });
745            });
746
747        let mut asr = vec![0.0f32; transpose.len()];
748        transpose
749            .par_iter()
750            .copied()
751            .zip(&mut asr)
752            .for_each(|(c, dest)| *dest = c.norm_sqr());
753
754        {
755            let mut asr =
756                BlurImageMut::borrow(&mut asr, image_width, image_height, FastBlurChannels::Plane);
757            stack_blur_f32(
758                &mut asr,
759                AnisotropicRadius::new(window_radius),
760                ThreadingPolicy::Adaptive,
761            )
762            .unwrap();
763        }
764
765        let max_asr = asr.par_iter().copied().reduce(|| f32::MIN, f32::max);
766        let min_asr = asr.par_iter().copied().reduce(|| f32::MAX, f32::min);
767        asr.par_iter_mut().for_each(|a| {
768            *a = (*a - min_asr) / (max_asr - min_asr).max(f32::EPSILON);
769        });
770
771        asr
772    };
773
774    let sr = {
775        let mut log_amplitude = vec![0.0f32; phase.len()];
776        amplitude
777            .into_par_iter()
778            .zip(&mut log_amplitude)
779            .for_each(|(a, dest)| *dest = a.max(f32::EPSILON).ln());
780
781        let mut log_blurred = log_amplitude.clone();
782
783        {
784            let mut amplitude = BlurImageMut::borrow(
785                &mut log_blurred,
786                image_width,
787                image_height,
788                FastBlurChannels::Plane,
789            );
790            stack_blur_f32(
791                &mut amplitude,
792                AnisotropicRadius::new(window_radius),
793                ThreadingPolicy::Adaptive,
794            )
795            .unwrap();
796        }
797
798        let mut residual = vec![0.0f32; image_width as usize * image_height as usize];
799        log_amplitude
800            .into_par_iter()
801            .zip(log_blurred)
802            .zip(&mut residual)
803            .for_each(|((a, b), dest)| {
804                *dest = a - b;
805            });
806
807        let sr_buffer = {
808            let ifft_row = planner.plan_fft_inverse(image_width as usize);
809            let ifft_col = planner.plan_fft_inverse(image_height as usize);
810
811            let mut buffer = vec![Complex::zero(); residual.len()];
812            residual
813                .into_par_iter()
814                .zip(phase.par_iter())
815                .zip(&mut buffer)
816                .for_each(|((a, p), dest)| {
817                    *dest = Complex {
818                        re: a.exp() * p.cos(),
819                        im: a.exp() * p.sin(),
820                    };
821                });
822
823            buffer
824                .par_chunks_mut(image_height as usize)
825                .for_each(|chunk| {
826                    ifft_col.process(chunk);
827                });
828
829            let mut transpose = vec![Complex::zero(); buffer.len()];
830            (0..image_width as usize)
831                .into_par_iter()
832                .zip(transpose.par_chunks_mut(image_height as usize))
833                .for_each(|(x, col)| {
834                    (0..image_height as usize)
835                        .into_par_iter()
836                        .zip(col)
837                        .for_each(|(y, dest)| {
838                            *dest = buffer[y * image_width as usize + x];
839                        });
840                });
841
842            transpose
843                .par_chunks_mut(image_height as usize)
844                .for_each(|chunk| {
845                    ifft_row.process(chunk);
846                });
847
848            (0..image_height as usize)
849                .into_par_iter()
850                .zip(buffer.par_chunks_mut(image_width as usize))
851                .for_each(|(y, row)| {
852                    (0..image_width as usize)
853                        .into_par_iter()
854                        .zip(row)
855                        .for_each(|(x, dest)| {
856                            *dest = transpose[x * image_height as usize + y];
857                        });
858                });
859
860            transpose
861        };
862
863        let mut saliency = vec![0.0f32; image_width as usize * image_height as usize];
864        sr_buffer
865            .into_par_iter()
866            .zip(&mut saliency)
867            .for_each(|(c, dest)| *dest = c.norm_sqr());
868
869        {
870            let mut saliency = BlurImageMut::borrow(
871                &mut saliency,
872                image_width,
873                image_height,
874                FastBlurChannels::Plane,
875            );
876            stack_blur_f32(
877                &mut saliency,
878                AnisotropicRadius::new(window_radius),
879                ThreadingPolicy::Adaptive,
880            )
881            .unwrap();
882        }
883
884        let max_saliency = saliency.par_iter().copied().reduce(|| f32::MIN, f32::max);
885        let min_saliency = saliency.par_iter().copied().reduce(|| f32::MAX, f32::min);
886        saliency.par_iter_mut().for_each(|s| {
887            *s = (*s - min_saliency) / (max_saliency - min_saliency).max(f32::EPSILON);
888        });
889
890        saliency
891    };
892
893    let msr = {
894        let mut log_amplitude = vec![0.0f32; phase.len()];
895        mod_amplitude
896            .into_par_iter()
897            .zip(&mut log_amplitude)
898            .for_each(|(a, dest)| *dest = a.max(f32::EPSILON).ln());
899
900        let mut log_blurred = log_amplitude.clone();
901
902        {
903            let mut amplitude = BlurImageMut::borrow(
904                &mut log_blurred,
905                image_width,
906                image_height,
907                FastBlurChannels::Plane,
908            );
909            stack_blur_f32(
910                &mut amplitude,
911                AnisotropicRadius::new(window_radius),
912                ThreadingPolicy::Adaptive,
913            )
914            .unwrap();
915        }
916
917        let mut residual = vec![0.0f32; image_width as usize * image_height as usize];
918        log_amplitude
919            .into_par_iter()
920            .zip(log_blurred)
921            .zip(&mut residual)
922            .for_each(|((a, b), dest)| {
923                *dest = a - b;
924            });
925
926        let mut mod_phase = vec![0.0f32; phase.len()];
927
928        let max_phase = phase.par_iter().copied().reduce(|| f32::MIN, f32::max);
929        let min_phase = phase.par_iter().copied().reduce(|| f32::MAX, f32::min);
930
931        phase
932            .into_par_iter()
933            .zip(&mut mod_phase)
934            .for_each(|(p, dest)| {
935                let x = ((p - min_phase) / (max_phase - min_phase).max(f32::EPSILON) - 0.5) * 2.0;
936                *dest = ((x * PI / 2.0).sin().powi(15) / 2.0 + 0.5) * (max_phase - min_phase)
937                    + min_phase;
938            });
939
940        let mut msr_buffer = vec![Complex::zero(); mod_phase.len()];
941        residual
942            .into_par_iter()
943            .zip(mod_phase)
944            .zip(&mut msr_buffer)
945            .for_each(|((a, p), dest)| {
946                *dest = Complex {
947                    re: a * p.cos(),
948                    im: a * p.sin(),
949                };
950            });
951
952        let ifft_row = planner.plan_fft_inverse(image_width as usize);
953        let ifft_col = planner.plan_fft_inverse(image_height as usize);
954        msr_buffer
955            .par_chunks_mut(image_height as usize)
956            .for_each(|chunk| {
957                ifft_col.process(chunk);
958            });
959
960        let mut transpose = vec![Complex::zero(); msr_buffer.len()];
961        (0..image_width as usize)
962            .into_par_iter()
963            .zip(transpose.par_chunks_mut(image_height as usize))
964            .for_each(|(x, col)| {
965                (0..image_height as usize)
966                    .into_par_iter()
967                    .zip(col)
968                    .for_each(|(y, dest)| {
969                        *dest = msr_buffer[y * image_width as usize + x];
970                    });
971            });
972
973        transpose
974            .par_chunks_mut(image_height as usize)
975            .for_each(|chunk| {
976                ifft_row.process(chunk);
977            });
978
979        (0..image_height as usize)
980            .into_par_iter()
981            .zip(msr_buffer.par_chunks_mut(image_width as usize))
982            .for_each(|(y, row)| {
983                (0..image_width as usize)
984                    .into_par_iter()
985                    .zip(row)
986                    .for_each(|(x, dest)| {
987                        *dest = transpose[x * image_height as usize + y];
988                    });
989            });
990
991        let mut msr = vec![0.0f32; transpose.len()];
992        transpose
993            .par_iter()
994            .copied()
995            .zip(&mut msr)
996            .for_each(|(c, dest)| *dest = c.norm_sqr());
997
998        {
999            let mut msr =
1000                BlurImageMut::borrow(&mut msr, image_width, image_height, FastBlurChannels::Plane);
1001            stack_blur_f32(
1002                &mut msr,
1003                AnisotropicRadius::new(window_radius),
1004                ThreadingPolicy::Adaptive,
1005            )
1006            .unwrap();
1007        }
1008
1009        let max_msr = msr.par_iter().copied().reduce(|| f32::MIN, f32::max);
1010        let min_msr = msr.par_iter().copied().reduce(|| f32::MAX, f32::min);
1011        msr.par_iter_mut().for_each(|s| {
1012            *s = (*s - min_msr) / (max_msr - min_msr).max(f32::EPSILON);
1013        });
1014
1015        msr
1016    };
1017
1018    Saliency {
1019        spectral_residual: sr,
1020        mod_spectral_residual: msr,
1021        phase_spectrum: pft,
1022        amplitude_spectrum: asr,
1023    }
1024}
1025
1026#[cfg(any(
1027    feature = "dump_l_saliency",
1028    feature = "dump_a_saliency",
1029    feature = "dump_b_saliency",
1030    feature = "dump_local_saliency",
1031    feature = "dump_weights"
1032))]
1033fn dump_intermediate(name: &str, buffer: &[u8], width: u32, height: u32) {
1034    use std::hash::{
1035        BuildHasher,
1036        Hasher,
1037        RandomState,
1038    };
1039    let rand = BuildHasher::build_hasher(&RandomState::new()).finish();
1040    image::save_buffer(
1041        format!("{name}-{rand}.png"),
1042        buffer,
1043        width,
1044        height,
1045        image::ColorType::L8,
1046    )
1047    .unwrap()
1048}