av_denoise/nlmeans/params.rs
1use super::{MotionCompensationMode, PrefilterMode};
2
3/// SSD normalisation reference, matching FFmpeg's nlmeans (255² for
4/// 8-bit normalisation). Distances are computed in `[0, 1]` units so
5/// this folds in the implied scale-up.
6pub(super) const NLM_NORM: f32 = 255.0 * 255.0;
7/// Legacy scaling factor inherited from FFmpeg's nlmeans; preserved so
8/// our `strength` parameter has equivalent meaning.
9pub(super) const NLM_LEGACY: f32 = 3.0;
10
11/// Patch radius threshold: above this the dispatcher switches to the
12/// separable path so the per-pixel cost stays linear in `patch_radius`.
13pub(super) const SEPARABLE_THRESHOLD: u32 = 8;
14
15/// Hard ceiling on `patch_radius`. The fused kernels load a
16/// `(block + 2·patch_radius)²` SMEM tile; values above this run out of
17/// SMEM on RDNA-class GPUs.
18pub const MAX_PATCH_RADIUS: u32 = 16;
19
20/// Hard ceiling on `search_radius`. The windowed kernel SMEM tile is
21/// `(block + 2·patch_radius + 2·search_radius)² × stored_ch × 4` bytes;
22/// the per-q dispatch path is also gated on this so launch counts stay
23/// sane (`(2·a+1)²` launches per frame).
24pub const MAX_SEARCH_RADIUS: u32 = 8;
25
26/// Hard ceiling on `temporal_radius`. The ring buffer is sized for
27/// `2·t + 1` frames; values above this consume excessive device memory
28/// (e.g. 1080p YUV at `t = 16` ≈ 540 MB just for input).
29pub const MAX_TEMPORAL_RADIUS: u32 = 8;
30
31#[derive(Debug, Clone, Copy, PartialEq, Eq)]
32/// How to apply denoising to the input frame channels.
33pub enum ChannelMode {
34 /// Single luminance channel. Distance scaled by 3.0.
35 Luma,
36 /// Two chroma channels (U, V). Distance scaled by 1.5.
37 Chroma,
38 /// Three channels (Y, U, V). Unscaled sum of squared differences.
39 Yuv,
40}
41
42impl ChannelMode {
43 /// Number of meaningful channels participating in distance/output.
44 pub fn count(self) -> u32 {
45 match self {
46 ChannelMode::Luma => 1,
47 ChannelMode::Chroma => 2,
48 ChannelMode::Yuv => 3,
49 }
50 }
51
52 /// Channels-per-pixel in GPU storage. Padded up to the next supported
53 /// vectorization factor so kernels can use coalesced `Line<f32>` reads
54 /// (backends only support power-of-two line sizes; YUV pads to 4).
55 pub fn storage_count(self) -> u32 {
56 match self {
57 ChannelMode::Luma => 1,
58 ChannelMode::Chroma => 2,
59 ChannelMode::Yuv => 4,
60 }
61 }
62}
63
64#[derive(Debug, Clone)]
65pub struct NlmParams {
66 /// Temporal radius. 0 = spatial only, d > 0 uses 2*d+1 frames.
67 pub temporal_radius: u32,
68 /// Search window half-size. Search window is (2*a+1)^2. Default: 2.
69 pub search_radius: u32,
70 /// Patch comparison half-size. Patch is (2*s+1)^2. Default: 4, range [0, 8].
71 pub patch_radius: u32,
72 /// Filtering strength. Higher = more smoothing. Default: 1.2.
73 pub strength: f32,
74 /// Self-weight multiplier. Default: 1.0. Set to 0 for pure NLM.
75 pub self_weight: f32,
76 /// Which channels to process.
77 pub channels: ChannelMode,
78 /// Reference clip source used for patch-distance / weight
79 /// computation. Default: `None`. When set, weights are derived
80 /// from a prefiltered or externally-supplied clip while pixel
81 /// accumulation continues to read the original input.
82 pub prefilter: PrefilterMode,
83 /// Motion-compensation mode. Default: `None`. When set to
84 /// `Mvtools`, each `denoise_submit` warps the temporal neighbours
85 /// into spatial alignment with the centre before NLM weighting.
86 /// Only takes effect when `temporal_radius > 0`.
87 pub motion_compensation: MotionCompensationMode,
88}
89
90impl Default for NlmParams {
91 fn default() -> Self {
92 Self {
93 temporal_radius: 0,
94 search_radius: 2,
95 patch_radius: 4,
96 strength: 1.2,
97 self_weight: 1.0,
98 channels: ChannelMode::Yuv,
99 prefilter: PrefilterMode::None,
100 motion_compensation: MotionCompensationMode::None,
101 }
102 }
103}
104
105impl NlmParams {
106 pub fn h2_inv_norm(&self) -> f32 {
107 let s_size = (2 * self.patch_radius + 1) * (2 * self.patch_radius + 1);
108 NLM_NORM / (NLM_LEGACY * self.strength * self.strength * s_size as f32)
109 }
110
111 pub(super) fn total_frames(&self) -> u32 {
112 1 + 2 * self.temporal_radius
113 }
114
115 /// Reject parameter combinations that would either fail to launch
116 /// (kernels hitting SMEM/register limits) or produce numerically
117 /// degenerate output. Called automatically by `NlmDenoiser::new`;
118 /// callers building params manually can invoke it directly to
119 /// surface errors before construction.
120 pub fn validate(&self) -> Result<(), anyhow::Error> {
121 if self.patch_radius > MAX_PATCH_RADIUS {
122 anyhow::bail!(
123 "patch_radius={} exceeds the supported maximum ({}); larger patches \
124 exhaust on-chip SMEM in the fused/windowed kernels",
125 self.patch_radius,
126 MAX_PATCH_RADIUS,
127 );
128 }
129
130 if self.search_radius > MAX_SEARCH_RADIUS {
131 anyhow::bail!(
132 "search_radius={} exceeds the supported maximum ({}); the windowed \
133 kernel allocates `(block + 2·patch_radius + 2·search_radius)²` of SMEM",
134 self.search_radius,
135 MAX_SEARCH_RADIUS,
136 );
137 }
138
139 if self.temporal_radius > MAX_TEMPORAL_RADIUS {
140 anyhow::bail!(
141 "temporal_radius={} exceeds the supported maximum ({}); the ring \
142 buffer grows linearly with the window size",
143 self.temporal_radius,
144 MAX_TEMPORAL_RADIUS,
145 );
146 }
147
148 if !(self.strength.is_finite() && self.strength > 0.0) {
149 anyhow::bail!(
150 "strength must be finite and > 0 (got {}); strength = 0 produces an \
151 infinite Welsch normalisation factor",
152 self.strength,
153 );
154 }
155
156 if !self.self_weight.is_finite() || self.self_weight < 0.0 {
157 anyhow::bail!("self_weight must be finite and >= 0 (got {})", self.self_weight,);
158 }
159
160 self.motion_compensation.validate()?;
161
162 Ok(())
163 }
164}