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 per-frame reference clip is produced.
9///
10/// `Bilateral` and any future GPU-internal variants run a kernel
11/// during `push_frame`. `External` requires the caller to supply a
12/// reference frame via [`super::NlmDenoiser::push_frame_with_reference`].
13/// `None` disables the reference path entirely (zero-cost).
14#[non_exhaustive]
15#[derive(Debug, Default, Clone, Copy, PartialEq)]
16pub enum PrefilterMode {
17    #[default]
18    None,
19    External,
20    Bilateral {
21        sigma_s: f32,
22        sigma_r: f32,
23    },
24    /// Spatial NLM pilot. Denoises each frame with the windowed
25    /// spatial kernel at push time and stores the result as the
26    /// reference clip, so patch distances are computed on a clean
27    /// image while accumulation still reads the noisy input.
28    NlmSpatial {
29        /// Multiplier on the main pass strength for the pilot pass.
30        strength_scale: f32,
31    },
32}
33
34/// Measured default for the pilot pass's relative strength, a
35/// multiplier on the main pass strength. A calibration sweep across
36/// noise levels found the XPSNR plateau optimum for
37/// `PrefilterMode::NlmSpatial` at this value.
38pub const DEFAULT_PILOT_STRENGTH_SCALE: f32 = 0.4;
39
40impl PrefilterMode {
41    /// Whether the denoiser needs to allocate the reference ring buffer.
42    pub(crate) fn needs_reference_buf(self) -> bool {
43        !matches!(self, Self::None)
44    }
45
46    /// Whether the variant computes its reference on the GPU during
47    /// `push_frame` (as opposed to consuming a caller-supplied clip).
48    pub(crate) fn is_gpu_internal(self) -> bool {
49        matches!(self, Self::Bilateral { .. } | Self::NlmSpatial { .. })
50    }
51}
52
53/// Inputs for a single-slot prefilter dispatch. Lives only for the
54/// duration of one `push_frame`, so borrows on the denoiser's buffers
55/// are sound.
56pub(crate) struct PrefilterCtx<'a> {
57    pub width: u32,
58    pub height: u32,
59    pub channels: u32,
60    pub stored_ch: u32,
61    pub frame_count: u32,
62    pub frame: u32,
63    pub input_buf: &'a Handle,
64    pub reference_buf: &'a Handle,
65}
66
67/// Dispatch the GPU prefilter for the most recently uploaded frame.
68/// `None` and `External` are no-ops.
69pub(crate) fn run_prefilter<R: Runtime>(
70    mode: PrefilterMode,
71    client: &ComputeClient<R>,
72    ctx: &PrefilterCtx<'_>,
73) -> Result<(), anyhow::Error> {
74    match mode {
75        PrefilterMode::None | PrefilterMode::External => Ok(()),
76        // The pilot needs the full accumulator context (accum,
77        // weight_sum, max_weight, h2_inv_norm), which `PrefilterCtx`
78        // doesn't carry, so `NlmDenoiser::run_nlm_spatial_pilot`
79        // dispatches it directly instead of going through this path.
80        PrefilterMode::NlmSpatial { .. } => Ok(()),
81        PrefilterMode::Bilateral { sigma_s, sigma_r } => {
82            bilateral::run_bilateral::<R>(client, ctx, sigma_s, sigma_r)
83        },
84    }
85}
86
87#[cfg(test)]
88mod tests {
89    use super::*;
90
91    #[test]
92    fn none_requires_no_reference_buffer() {
93        assert!(!PrefilterMode::None.needs_reference_buf());
94        assert!(!PrefilterMode::None.is_gpu_internal());
95    }
96
97    #[test]
98    fn external_needs_buffer_but_not_gpu() {
99        assert!(PrefilterMode::External.needs_reference_buf());
100        assert!(!PrefilterMode::External.is_gpu_internal());
101    }
102
103    #[test]
104    fn bilateral_is_gpu_internal() {
105        let m = PrefilterMode::Bilateral {
106            sigma_s: 3.0,
107            sigma_r: 0.02,
108        };
109
110        assert!(m.needs_reference_buf());
111        assert!(m.is_gpu_internal());
112    }
113
114    #[test]
115    fn nlm_spatial_is_gpu_internal() {
116        let m = PrefilterMode::NlmSpatial { strength_scale: 1.0 };
117
118        assert!(m.needs_reference_buf());
119        assert!(m.is_gpu_internal());
120    }
121
122    #[test]
123    fn bilateral_radius_truncates_at_two_sigma() {
124        assert_eq!(bilateral_radius(0.1), 1);
125        assert_eq!(bilateral_radius(1.0), 2);
126        assert_eq!(bilateral_radius(3.0), 6);
127        assert_eq!(bilateral_radius(3.5), 7);
128    }
129}
av_denoise/nlmeans/prefilter/mod.rs

av_denoise/nlmeans/prefilter/
mod.rs
