av_denoise/nlmeans/prefilter/
mod.rs

1mod bilateral;
2
3pub use bilateral::bilateral_radius;
4pub(crate) use bilateral::inv_two_sigma_sq;
5use cubecl::prelude::*;
6use cubecl::server::Handle;
7
8/// How the reference image for each frame is produced.
9///
10/// NLM compares patches to decide how much two pixels look alike. Doing
11/// that on a noisy image means comparing noise as well as content, so a
12/// cleaner reference image can give better weights.
13///
14/// The pixels being averaged always come from the original input. Only
15/// the weights change.
16#[non_exhaustive]
17#[derive(Debug, Default, Clone, Copy, PartialEq)]
18pub enum PrefilterMode {
19    /// No reference image, so patches are compared on the noisy input.
20    /// This costs nothing extra.
21    #[default]
22    None,
23    /// The caller supplies the reference frame through
24    /// [`super::NlmDenoiser::push_frame_with_reference`].
25    External,
26    /// A quick bilateral blur run on the GPU at push time.
27    Bilateral { sigma_s: f32, sigma_r: f32 },
28    /// A spatial NLM pilot pass.
29    ///
30    /// Each frame is denoised with the windowed spatial kernel at push
31    /// time and the result is kept as the reference image.
32    NlmSpatial {
33        /// How much of the main pass strength the pilot pass uses.
34        strength_scale: f32,
35    },
36}
37
38/// The measured default strength for the pilot pass, as a multiplier on
39/// the main pass strength.
40///
41/// A calibration sweep across noise levels put the XPSNR plateau for
42/// `PrefilterMode::NlmSpatial` at this value.
43pub const DEFAULT_PILOT_STRENGTH_SCALE: f32 = 0.4;
44
45impl PrefilterMode {
46    /// Whether the denoiser needs to allocate the reference ring buffer.
47    pub(crate) fn needs_reference_buf(self) -> bool {
48        !matches!(self, Self::None)
49    }
50
51    /// Whether this mode builds its reference on the GPU during
52    /// `push_frame`, rather than taking one from the caller.
53    pub(crate) fn is_gpu_internal(self) -> bool {
54        matches!(self, Self::Bilateral { .. } | Self::NlmSpatial { .. })
55    }
56}
57
58/// The inputs one prefilter dispatch needs.
59///
60/// This lives only for the length of a single `push_frame`, which is
61/// what makes the borrows on the denoiser's buffers sound.
62pub(crate) struct PrefilterCtx<'a> {
63    pub width: u32,
64    pub height: u32,
65    pub channels: u32,
66    pub stored_ch: u32,
67    pub frame_count: u32,
68    pub frame: u32,
69    pub input_buf: &'a Handle,
70    pub reference_buf: &'a Handle,
71}
72
73/// Runs the GPU prefilter for the frame that was uploaded last.
74///
75/// `None` and `External` do nothing here.
76pub(crate) fn run_prefilter<R: Runtime>(
77    mode: PrefilterMode,
78    client: &ComputeClient<R>,
79    ctx: &PrefilterCtx<'_>,
80) -> Result<(), anyhow::Error> {
81    match mode {
82        PrefilterMode::None | PrefilterMode::External => Ok(()),
83        // The pilot needs the full accumulator context, meaning accum,
84        // weight_sum, max_weight, and h2_inv_norm, which `PrefilterCtx`
85        // does not carry. `NlmDenoiser::run_nlm_spatial_pilot`
86        // dispatches it directly instead of coming through here.
87        PrefilterMode::NlmSpatial { .. } => Ok(()),
88        PrefilterMode::Bilateral { sigma_s, sigma_r } => {
89            bilateral::run_bilateral::<R>(client, ctx, sigma_s, sigma_r)
90        },
91    }
92}
93
94#[cfg(test)]
95mod tests {
96    use super::*;
97
98    #[test]
99    fn none_requires_no_reference_buffer() {
100        assert!(!PrefilterMode::None.needs_reference_buf());
101        assert!(!PrefilterMode::None.is_gpu_internal());
102    }
103
104    #[test]
105    fn external_needs_buffer_but_not_gpu() {
106        assert!(PrefilterMode::External.needs_reference_buf());
107        assert!(!PrefilterMode::External.is_gpu_internal());
108    }
109
110    #[test]
111    fn bilateral_is_gpu_internal() {
112        let m = PrefilterMode::Bilateral {
113            sigma_s: 3.0,
114            sigma_r: 0.02,
115        };
116
117        assert!(m.needs_reference_buf());
118        assert!(m.is_gpu_internal());
119    }
120
121    #[test]
122    fn nlm_spatial_is_gpu_internal() {
123        let m = PrefilterMode::NlmSpatial { strength_scale: 1.0 };
124
125        assert!(m.needs_reference_buf());
126        assert!(m.is_gpu_internal());
127    }
128
129    #[test]
130    fn bilateral_radius_truncates_at_two_sigma() {
131        assert_eq!(bilateral_radius(0.1), 1);
132        assert_eq!(bilateral_radius(1.0), 2);
133        assert_eq!(bilateral_radius(3.0), 6);
134        assert_eq!(bilateral_radius(3.5), 7);
135    }
136}
av_denoise/nlmeans/prefilter/mod.rs

av_denoise/nlmeans/prefilter/
mod.rs
