av_denoise/nlmeans/
mod.rs

1//! The non-local means denoiser that sits behind [`crate::Denoiser`].
2//!
3//! Non-local means cleans a pixel by finding patches elsewhere that look
4//! like the patch around it, then averaging them. Similar patches get a
5//! large weight and dissimilar ones get almost none, so flat areas
6//! smooth out while edges survive.
7//!
8//! The search can reach across neighbouring frames as well as within one
9//! frame, which is what the temporal radius controls.
10//!
11//! # Layout
12//!
13//! `params` holds the tuning values and the calibrated defaults, and
14//! [`NlmParams`] is the single struct everything else is built from.
15//!
16//! [`NlmDenoiser`] owns the GPU buffers and the frame ring, and
17//! `dispatch` turns one set of parameters into the sequence of kernel
18//! launches that produces a frame.
19//!
20//! [`kernels`] holds the GPU code itself. `noise` measures how noisy a
21//! frame is, [`motion`] tracks movement between frames, and
22//! [`prefilter`] builds the cleaner reference image that patches are
23//! compared against.
24
25pub mod kernels;
26pub mod motion;
27pub mod prefilter;
28
29mod align;
30mod denoiser;
31mod dispatch;
32mod noise;
33mod params;
34mod pending;
35
36// Every test in this tree runs against a real GPU runtime, see
37// `tests::helpers::R`, so it only builds when a wgpu-backed feature is
38// enabled. A cpu-only build skips it entirely, and the
39// `cpu_smoke_tests` module in `src/denoiser.rs` covers that backend
40// instead.
41#[cfg(all(test, any(feature = "vulkan", feature = "metal")))]
42mod tests;
43
44pub(crate) use denoiser::RingView;
45pub use denoiser::{GpuOutput, NlmDenoiser};
46pub use motion::{MotionCompensationMode, MotionEstimation, MotionSearch};
47pub use params::{
48    ChannelMode,
49    HqParams,
50    MAX_PATCH_RADIUS,
51    MAX_SEARCH_RADIUS,
52    MAX_TEMPORAL_RADIUS,
53    MIN_FRAME_DIM,
54    NlmParams,
55    hq_default_strength,
56    validate_dimensions,
57};
58pub use pending::Pending;
59pub use prefilter::{DEFAULT_PILOT_STRENGTH_SCALE, PrefilterMode};
60
61/// Cube X dimension for tile-heavy fused/separable kernels.
62pub const BLOCK_X: u32 = 32;
63/// Cube Y dimension for tile-heavy fused/separable kernels.
64pub const BLOCK_Y: u32 = 8;
65
66/// Cube shape for the per-pixel `nlm_accumulate` kernel, which has no
67/// shared-memory tile.
68///
69/// On RDNA-class GPUs this shape benchmarks 10 to 25% faster than the
70/// tile-heavy default. The kernel waits on memory rather than compute,
71/// so the extra threads hide the load latency.
72pub const BLOCK_X_THIN: u32 = 32;
73pub const BLOCK_Y_THIN: u32 = 16;
74
75/// Largest 1D grid a dispatch may ask for, set by the WebGPU and Vulkan
76/// limits.
77pub(crate) const MAX_GRID_1D: u32 = 65535;
78
79/// Block size for 1D utility kernels (copy, zero).
80pub(crate) const BLOCK_1D: u32 = 256;
81
82/// Bit depth of a source's samples.
83///
84/// Normalisation divides by [`Depth::max_value`], so a value in
85/// normalised units means the same thing at every depth.
86#[derive(Debug, Clone, Copy, PartialEq, Eq)]
87pub enum Depth {
88    Eight,
89    Ten,
90    Twelve,
91}
92
93/// Returned when a source declares a bit depth the denoiser does not handle.
94#[derive(Debug, thiserror::Error)]
95#[error("unsupported bit depth {0}, av-denoise supports 8, 10, and 12-bit")]
96pub struct UnsupportedDepthError(pub usize);
97
98impl Depth {
99    /// Maps a declared bit depth onto a [`Depth`].
100    pub fn from_bits(bits: usize) -> Result<Self, UnsupportedDepthError> {
101        match bits {
102            8 => Ok(Depth::Eight),
103            10 => Ok(Depth::Ten),
104            12 => Ok(Depth::Twelve),
105            other => Err(UnsupportedDepthError(other)),
106        }
107    }
108
109    /// Bits per sample.
110    pub fn bits(self) -> usize {
111        match self {
112            Depth::Eight => 8,
113            Depth::Ten => 10,
114            Depth::Twelve => 12,
115        }
116    }
117
118    /// Bytes each sample takes up on the wire.
119    ///
120    /// Depths above 8 use a little-endian 16-bit word.
121    pub fn bytes_per_sample(self) -> usize {
122        match self {
123            Depth::Eight => 1,
124            Depth::Ten | Depth::Twelve => 2,
125        }
126    }
127
128    /// The largest sample value this depth can hold, which is also the
129    /// normalisation divisor.
130    pub fn max_value(self) -> f32 {
131        ((1u32 << self.bits()) - 1) as f32
132    }
133
134    /// The sample value that means neutral chroma at this depth.
135    pub fn neutral_chroma(self) -> u16 {
136        1 << (self.bits() - 1)
137    }
138}
139
140/// Scales native-depth samples into normalised `[0, 1]` f32.
141pub fn normalize(input: &[u16], depth: Depth) -> Vec<f32> {
142    let max = depth.max_value();
143    input.iter().map(|&v| v as f32 / max).collect()
144}
145
146/// Reverse of [`normalize`].
147///
148/// Values outside `[0, 1]` are clamped, and `NaN` becomes 0.
149pub fn denormalize(input: &[f32], depth: Depth) -> Vec<u16> {
150    let max = depth.max_value();
151    input
152        .iter()
153        .map(|&v| (v * max).round().clamp(0.0, max) as u16)
154        .collect()
155}
156
157#[cfg(test)]
158mod depth_tests {
159    use super::*;
160
161    #[test]
162    fn from_bits_accepts_supported_depths() {
163        assert_eq!(Depth::from_bits(8).unwrap(), Depth::Eight);
164        assert_eq!(Depth::from_bits(10).unwrap(), Depth::Ten);
165        assert_eq!(Depth::from_bits(12).unwrap(), Depth::Twelve);
166    }
167
168    #[test]
169    fn from_bits_rejects_unsupported_depths() {
170        for bits in [0, 9, 14, 16] {
171            let err = Depth::from_bits(bits).expect_err("expected rejection");
172            assert!(
173                err.to_string().contains(&bits.to_string()),
174                "error should name the depth, got {err}"
175            );
176        }
177    }
178
179    #[test]
180    fn depth_properties_match_the_format() {
181        assert_eq!(Depth::Eight.bytes_per_sample(), 1);
182        assert_eq!(Depth::Ten.bytes_per_sample(), 2);
183        assert_eq!(Depth::Twelve.bytes_per_sample(), 2);
184
185        assert_eq!(Depth::Eight.max_value(), 255.0);
186        assert_eq!(Depth::Ten.max_value(), 1023.0);
187        assert_eq!(Depth::Twelve.max_value(), 4095.0);
188
189        assert_eq!(Depth::Eight.neutral_chroma(), 128);
190        assert_eq!(Depth::Ten.neutral_chroma(), 512);
191        assert_eq!(Depth::Twelve.neutral_chroma(), 2048);
192    }
193
194    /// Limited-range black and white land on matching normalised values
195    /// at every depth, which is what lets every calibrated constant in
196    /// the library stay depth-independent.
197    ///
198    /// The match is within one 8-bit code level rather than exact. ITU
199    /// defines the limited-range endpoints as exact multiples, so 235
200    /// becomes 940 and then 3760, but full scale is not a multiple,
201    /// because 255 becomes 1023 and then 4095.
202    ///
203    /// That leaves 235/255 and 940/1023 differing by 0.0027, roughly
204    /// 0.69 of an 8-bit step. Agreement below one step is the real
205    /// property here.
206    #[test]
207    fn normalized_scale_is_identical_across_depths() {
208        /// One 8-bit code level, the precision the endpoints agree to.
209        const TOL: f32 = 1.0 / 255.0;
210
211        let eight = normalize(&[16, 235], Depth::Eight);
212        let ten = normalize(&[64, 940], Depth::Ten);
213        let twelve = normalize(&[256, 3760], Depth::Twelve);
214
215        for (a, b) in eight.iter().zip(ten.iter()) {
216            assert!((a - b).abs() < TOL, "8-bit {a} vs 10-bit {b}");
217        }
218        for (a, b) in eight.iter().zip(twelve.iter()) {
219            assert!((a - b).abs() < TOL, "8-bit {a} vs 12-bit {b}");
220        }
221    }
222
223    #[test]
224    fn normalization_round_trips_at_every_depth() {
225        for depth in [Depth::Eight, Depth::Ten, Depth::Twelve] {
226            let max = depth.max_value() as u16;
227            let original: Vec<u16> = vec![0, 1, 16, 64, 128, 235, max / 2, max - 1, max];
228            let restored = denormalize(&normalize(&original, depth), depth);
229            assert_eq!(original, restored, "round trip failed at {depth:?}");
230        }
231    }
232
233    #[test]
234    fn denormalize_clamps_out_of_range_input() {
235        let out = denormalize(&[-0.5, 0.0, 1.0, 1.5], Depth::Ten);
236        assert_eq!(out, vec![0, 0, 1023, 1023]);
237    }
238}
av_denoise/nlmeans/mod.rs

av_denoise/nlmeans/
mod.rs
