av_denoise_core/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(crate) use pending::start_readback;
59pub use pending::{Pending, TryWait};
60pub use prefilter::{DEFAULT_PILOT_STRENGTH_SCALE, PrefilterMode, parse_prefilter};
61
62/// Cube X dimension for tile-heavy fused/separable kernels.
63pub const BLOCK_X: u32 = 32;
64/// Cube Y dimension for tile-heavy fused/separable kernels.
65pub const BLOCK_Y: u32 = 8;
66
67/// Cube shape for the per-pixel `nlm_accumulate` kernel, which has no
68/// shared-memory tile.
69///
70/// On RDNA-class GPUs this shape benchmarks 10 to 25% faster than the
71/// tile-heavy default. The kernel waits on memory rather than compute,
72/// so the extra threads hide the load latency.
73pub const BLOCK_X_THIN: u32 = 32;
74pub const BLOCK_Y_THIN: u32 = 16;
75
76/// Largest 1D grid a dispatch may ask for, set by the WebGPU and Vulkan
77/// limits.
78pub(crate) const MAX_GRID_1D: u32 = 65535;
79
80/// Block size for 1D utility kernels (copy, zero).
81pub(crate) const BLOCK_1D: u32 = 256;
82
83/// Bit depth of a source's samples.
84///
85/// Normalisation divides by [`Depth::max_value`], so a value in
86/// normalised units means the same thing at every depth.
87#[derive(Debug, Clone, Copy, PartialEq, Eq)]
88pub enum Depth {
89    Eight,
90    Ten,
91    Twelve,
92}
93
94/// Returned when a source declares a bit depth the denoiser does not handle.
95#[derive(Debug, thiserror::Error)]
96#[error("unsupported bit depth {0}, av-denoise supports 8, 10, and 12-bit")]
97pub struct UnsupportedDepthError(pub usize);
98
99impl Depth {
100    /// Maps a declared bit depth onto a [`Depth`].
101    pub fn from_bits(bits: usize) -> Result<Self, UnsupportedDepthError> {
102        match bits {
103            8 => Ok(Depth::Eight),
104            10 => Ok(Depth::Ten),
105            12 => Ok(Depth::Twelve),
106            other => Err(UnsupportedDepthError(other)),
107        }
108    }
109
110    /// Bits per sample.
111    pub fn bits(self) -> usize {
112        match self {
113            Depth::Eight => 8,
114            Depth::Ten => 10,
115            Depth::Twelve => 12,
116        }
117    }
118
119    /// Bytes each sample takes up on the wire.
120    ///
121    /// Depths above 8 use a little-endian 16-bit word.
122    pub fn bytes_per_sample(self) -> usize {
123        match self {
124            Depth::Eight => 1,
125            Depth::Ten | Depth::Twelve => 2,
126        }
127    }
128
129    /// The largest sample value this depth can hold, which is also the
130    /// normalisation divisor.
131    pub fn max_value(self) -> f32 {
132        ((1u32 << self.bits()) - 1) as f32
133    }
134
135    /// The sample value that means neutral chroma at this depth.
136    pub fn neutral_chroma(self) -> u16 {
137        1 << (self.bits() - 1)
138    }
139
140    /// How [`crate::nlmeans::kernels::gpu_pack_wire`] packs samples at
141    /// this depth.
142    pub fn wire_pack(self) -> WirePack {
143        WirePack {
144            max: self.max_value(),
145            samples_per_word: 4 / self.bytes_per_sample() as u32,
146        }
147    }
148}
149
150/// The quantisation scale and lane count `gpu_pack_wire` packs one
151/// sample with.
152///
153/// The kernel shifts a quantised sample by `lane * (32 / samples_per_word)`
154/// and never masks it, so a `max` wider than the lane holds spills bits
155/// into the neighbouring sample. Both values come from one [`Depth`]
156/// through [`Depth::wire_pack`], which makes that pairing impossible to
157/// get wrong.
158#[derive(Debug, Clone, Copy, PartialEq)]
159pub struct WirePack {
160    max: f32,
161    samples_per_word: u32,
162}
163
164impl WirePack {
165    /// The scale a normalised value is quantised against.
166    pub fn max(self) -> f32 {
167        self.max
168    }
169
170    /// How many samples share one `u32` word.
171    pub fn samples_per_word(self) -> u32 {
172        self.samples_per_word
173    }
174}
175
176/// Scales native-depth samples into normalised `[0, 1]` f32.
177pub fn normalize(input: &[u16], depth: Depth) -> Vec<f32> {
178    let max = depth.max_value();
179    input.iter().map(|&v| v as f32 / max).collect()
180}
181
182/// Reverse of [`normalize`].
183///
184/// Values outside `[0, 1]` are clamped, and `NaN` becomes 0.
185pub fn denormalize(input: &[f32], depth: Depth) -> Vec<u16> {
186    let max = depth.max_value();
187    input
188        .iter()
189        .map(|&v| (v * max).round().clamp(0.0, max) as u16)
190        .collect()
191}
192
193#[cfg(test)]
194mod depth_tests {
195    use super::*;
196
197    #[test]
198    fn from_bits_accepts_supported_depths() {
199        assert_eq!(Depth::from_bits(8).unwrap(), Depth::Eight);
200        assert_eq!(Depth::from_bits(10).unwrap(), Depth::Ten);
201        assert_eq!(Depth::from_bits(12).unwrap(), Depth::Twelve);
202    }
203
204    #[test]
205    fn from_bits_rejects_unsupported_depths() {
206        for bits in [0, 9, 14, 16] {
207            let err = Depth::from_bits(bits).expect_err("expected rejection");
208            assert!(
209                err.to_string().contains(&bits.to_string()),
210                "error should name the depth, got {err}"
211            );
212        }
213    }
214
215    #[test]
216    fn depth_properties_match_the_format() {
217        assert_eq!(Depth::Eight.bytes_per_sample(), 1);
218        assert_eq!(Depth::Ten.bytes_per_sample(), 2);
219        assert_eq!(Depth::Twelve.bytes_per_sample(), 2);
220
221        assert_eq!(Depth::Eight.max_value(), 255.0);
222        assert_eq!(Depth::Ten.max_value(), 1023.0);
223        assert_eq!(Depth::Twelve.max_value(), 4095.0);
224
225        assert_eq!(Depth::Eight.neutral_chroma(), 128);
226        assert_eq!(Depth::Ten.neutral_chroma(), 512);
227        assert_eq!(Depth::Twelve.neutral_chroma(), 2048);
228    }
229
230    /// Limited-range black and white land on matching normalised values
231    /// at every depth, which is what lets every calibrated constant in
232    /// the library stay depth-independent.
233    ///
234    /// The match is within one 8-bit code level rather than exact. ITU
235    /// defines the limited-range endpoints as exact multiples, so 235
236    /// becomes 940 and then 3760, but full scale is not a multiple,
237    /// because 255 becomes 1023 and then 4095.
238    ///
239    /// That leaves 235/255 and 940/1023 differing by 0.0027, roughly
240    /// 0.69 of an 8-bit step. Agreement below one step is the real
241    /// property here.
242    #[test]
243    fn normalized_scale_is_identical_across_depths() {
244        /// One 8-bit code level, the precision the endpoints agree to.
245        const TOL: f32 = 1.0 / 255.0;
246
247        let eight = normalize(&[16, 235], Depth::Eight);
248        let ten = normalize(&[64, 940], Depth::Ten);
249        let twelve = normalize(&[256, 3760], Depth::Twelve);
250
251        for (a, b) in eight.iter().zip(ten.iter()) {
252            assert!((a - b).abs() < TOL, "8-bit {a} vs 10-bit {b}");
253        }
254        for (a, b) in eight.iter().zip(twelve.iter()) {
255            assert!((a - b).abs() < TOL, "8-bit {a} vs 12-bit {b}");
256        }
257    }
258
259    #[test]
260    fn normalization_round_trips_at_every_depth() {
261        for depth in [Depth::Eight, Depth::Ten, Depth::Twelve] {
262            let max = depth.max_value() as u16;
263            let original: Vec<u16> = vec![0, 1, 16, 64, 128, 235, max / 2, max - 1, max];
264            let restored = denormalize(&normalize(&original, depth), depth);
265            assert_eq!(original, restored, "round trip failed at {depth:?}");
266        }
267    }
268
269    #[test]
270    fn denormalize_clamps_out_of_range_input() {
271        let out = denormalize(&[-0.5, 0.0, 1.0, 1.5], Depth::Ten);
272        assert_eq!(out, vec![0, 0, 1023, 1023]);
273    }
274}
av_denoise_core/nlmeans/mod.rs

av_denoise_core/nlmeans/
mod.rs
