av_denoise/nlmeans/prefilter/
mod.rs

1mod bilateral;
2
3pub use bilateral::bilateral_radius;
4use cubecl::prelude::*;
5use cubecl::server::Handle;
6
7/// How the per-frame reference clip is produced.
8///
9/// `Bilateral` and any future GPU-internal variants run a kernel
10/// during `push_frame`. `External` requires the caller to supply a
11/// reference frame via [`super::NlmDenoiser::push_frame_with_reference`].
12/// `None` disables the reference path entirely (zero-cost).
13#[non_exhaustive]
14#[derive(Debug, Default, Clone, Copy, PartialEq)]
15pub enum PrefilterMode {
16    #[default]
17    None,
18    External,
19    Bilateral {
20        sigma_s: f32,
21        sigma_r: f32,
22    },
23}
24
25impl PrefilterMode {
26    /// Whether the denoiser needs to allocate the reference ring buffer.
27    pub(crate) fn needs_reference_buf(self) -> bool {
28        !matches!(self, Self::None)
29    }
30
31    /// Whether the variant computes its reference on the GPU during
32    /// `push_frame` (as opposed to consuming a caller-supplied clip).
33    pub(crate) fn is_gpu_internal(self) -> bool {
34        matches!(self, Self::Bilateral { .. })
35    }
36}
37
38/// Inputs for a single-slot prefilter dispatch. Lives only for the
39/// duration of one `push_frame`, so borrows on the denoiser's buffers
40/// are sound.
41pub(crate) struct PrefilterCtx<'a> {
42    pub width: u32,
43    pub height: u32,
44    pub channels: u32,
45    pub stored_ch: u32,
46    pub frame_count: u32,
47    pub frame: u32,
48    pub input_buf: &'a Handle,
49    pub reference_buf: &'a Handle,
50}
51
52/// Dispatch the GPU prefilter for the most recently uploaded frame.
53/// `None` and `External` are no-ops.
54pub(crate) fn run_prefilter<R: Runtime>(
55    mode: PrefilterMode,
56    client: &ComputeClient<R>,
57    ctx: &PrefilterCtx<'_>,
58) -> Result<(), anyhow::Error> {
59    match mode {
60        PrefilterMode::None | PrefilterMode::External => Ok(()),
61        PrefilterMode::Bilateral { sigma_s, sigma_r } => {
62            bilateral::run_bilateral::<R>(client, ctx, sigma_s, sigma_r)
63        },
64    }
65}
66
67#[cfg(test)]
68mod tests {
69    use super::*;
70
71    #[test]
72    fn none_requires_no_reference_buffer() {
73        assert!(!PrefilterMode::None.needs_reference_buf());
74        assert!(!PrefilterMode::None.is_gpu_internal());
75    }
76
77    #[test]
78    fn external_needs_buffer_but_not_gpu() {
79        assert!(PrefilterMode::External.needs_reference_buf());
80        assert!(!PrefilterMode::External.is_gpu_internal());
81    }
82
83    #[test]
84    fn bilateral_is_gpu_internal() {
85        let m = PrefilterMode::Bilateral {
86            sigma_s: 3.0,
87            sigma_r: 0.02,
88        };
89
90        assert!(m.needs_reference_buf());
91        assert!(m.is_gpu_internal());
92    }
93
94    #[test]
95    fn bilateral_radius_truncates_at_two_sigma() {
96        assert_eq!(bilateral_radius(0.1), 1);
97        assert_eq!(bilateral_radius(1.0), 2);
98        assert_eq!(bilateral_radius(3.0), 6);
99        assert_eq!(bilateral_radius(3.5), 7);
100    }
101}
av_denoise/nlmeans/prefilter/mod.rs

av_denoise/nlmeans/prefilter/
mod.rs
