Skip to main content

av_denoise_vs/
params.rs

1//! Turns a VapourSynth clip's format and a filter's script arguments
2//! into the option types `av-denoise-core` denoises with.
3//!
4//! Everything here is a pure function over plain values, with no
5//! VapourSynth core and no GPU, so the whole accept/reject matrix is
6//! unit-testable. [`Format`](vapoursynth::format::Format) itself cannot
7//! be built outside a running core, so [`layout_from_format`] takes a
8//! [`RawFormat`] of the plain fields it needs instead. The caller in
9//! `filter.rs` does the short extraction from a real `Format`.
10
11use av_denoise_core::accelerate::{Accelerator, get_default_accelerators};
12use av_denoise_core::{
13    Algorithm,
14    ChannelIntent,
15    DenoisingMode,
16    Depth,
17    Device,
18    FrameLayout,
19    HqParams,
20    MotionCompensationMode,
21    MotionSearch,
22    Nl4dOptions,
23    NlmTuning,
24    NlmeansHqOptions,
25    NlmeansOptions,
26    NlmeansVariant,
27    PlaneOptions,
28    PrefilterMode,
29    Preset,
30    Subsampling,
31    nl4d_spatial_radius_for,
32    nl4d_temporal_radius_for,
33    nlmeans_search_radius_for,
34    nlmeans_temporal_radius_for,
35    nlmeans_variant_for,
36    parse_prefilter,
37};
38use vapoursynth::format::{ColorFamily, SampleType};
39
40/// The handful of format fields [`layout_from_format`] actually reads.
41///
42/// The real caller is `vapoursynth::format::Format`, which wraps a
43/// pointer only a running VapourSynth core can hand out, so it cannot be
44/// built in a unit test. A caller with a real `Format` builds one of
45/// these from `format.sample_type()`, `format.bits_per_sample()`,
46/// `format.sub_sampling_w()`, `format.sub_sampling_h()`, and
47/// `format.color_family()`.
48#[derive(Debug, Clone, Copy)]
49pub struct RawFormat {
50    pub sample_type: SampleType,
51    pub bits_per_sample: u8,
52    pub subsampling_w: u8,
53    pub subsampling_h: u8,
54    pub color_family: ColorFamily,
55}
56
57/// Validates a clip's format and turns it into a [`FrameLayout`].
58///
59/// Accepts integer YUV420, YUV422, and YUV444 sources at 8, 10, or
60/// 12-bit. Rejects RGB, since the denoiser's channel distance weights
61/// are calibrated for YUV. Rejects float sample types and any other
62/// chroma subsampling.
63///
64/// Rejects GRAY too. Core's [`Subsampling`] has no "no chroma" variant,
65/// so a GRAY source would have to be represented as YUV444, which makes
66/// [`av_denoise_core::frame::FrameLayout::chroma_dims`] report
67/// full-resolution chroma planes that do not exist. The filter would
68/// then have to fabricate and push full-size neutral chroma every
69/// frame, four times the real data volume of true 4:2:0 chroma, purely
70/// to work around a gap in the geometry type. GRAY is out of scope
71/// until core can represent a source with no chroma planes at all.
72pub fn layout_from_format(format: RawFormat, width: u32, height: u32) -> Result<FrameLayout, anyhow::Error> {
73    match format.color_family {
74        ColorFamily::YUV => {},
75        ColorFamily::Gray => {
76            anyhow::bail!(
77                "GRAY clips are not supported, av-denoise-vs only accepts YUV420, YUV422, and YUV444 sources. Convert the input to YUV first, for example with `ffmpeg -pix_fmt yuv420p`"
78            );
79        },
80        other => {
81            anyhow::bail!(
82                "{other:?} clips are not supported, av-denoise's channel distance weights are calibrated for YUV. Convert the input to YUV first, for example with `ffmpeg -pix_fmt yuv420p`"
83            );
84        },
85    }
86
87    if format.sample_type == SampleType::Float {
88        anyhow::bail!(
89            "float sample types are not supported, av-denoise expects integer YUV samples. Convert to an integer format first, for example with `ffmpeg -pix_fmt yuv420p`"
90        );
91    }
92
93    let depth = Depth::from_bits(format.bits_per_sample as usize)?;
94
95    let subsampling = match (format.subsampling_w, format.subsampling_h) {
96        (0, 0) => Subsampling::Yuv444,
97        (1, 0) => Subsampling::Yuv422,
98        (1, 1) => Subsampling::Yuv420,
99        (w, h) => {
100            anyhow::bail!(
101                "unsupported chroma subsampling (subsampling_w={w}, subsampling_h={h}), av-denoise-vs accepts YUV420, YUV422, and YUV444"
102            );
103        },
104    };
105
106    Ok(FrameLayout {
107        width,
108        height,
109        subsampling,
110        depth,
111    })
112}
113
114/// Which denoising algorithm a filter function runs.
115///
116/// `avd.Nlmeans` builds [`AlgorithmKind::Nlmeans`], `avd.Nl4d` builds
117/// [`AlgorithmKind::Nl4d`]. This has no `Hq` variant since the
118/// VapourSynth plugin does not expose the HQ variant separately, it
119/// takes a plain algorithm choice per filter function.
120#[derive(Debug, Clone, Copy, PartialEq, Eq)]
121pub enum AlgorithmKind {
122    Nlmeans,
123    Nl4d,
124}
125
126/// The name a [`NlmeansVariant`] parses back from, used in error
127/// messages.
128fn variant_name(variant: NlmeansVariant) -> &'static str {
129    match variant {
130        NlmeansVariant::Fast => "fast",
131        NlmeansVariant::Hq => "hq",
132    }
133}
134
135/// Resolves an explicit `variant` string into an [`NlmeansVariant`].
136///
137/// Uses [`av_denoise_core`]'s own parser, the same one the CLI's
138/// `--variant` flag resolves through, so a name accepted on the CLI is
139/// accepted here too.
140fn parse_variant(raw: &str) -> Result<NlmeansVariant, anyhow::Error> {
141    raw.parse::<NlmeansVariant>()
142        .map_err(|_| anyhow::anyhow!("unknown variant '{raw}', expected one of fast, hq"))
143}
144
145/// Resolves an explicit `preset` string into a [`Preset`].
146///
147/// Uses [`av_denoise_core`]'s own parser, the same one the CLI's
148/// `--preset` flag resolves through, so a name accepted on the CLI is
149/// accepted here too.
150fn parse_preset(raw: &str) -> Result<Preset, anyhow::Error> {
151    raw.parse::<Preset>().map_err(|_| {
152        anyhow::anyhow!("unknown preset '{raw}', expected one of veryfast, fast, base, slow, veryslow")
153    })
154}
155
156/// The raw script arguments a filter function receives, before they are
157/// validated and folded into a [`PlaneOptions`].
158///
159/// Every field is optional. An unset field falls back to the library's
160/// own default for whichever algorithm is being built.
161#[derive(Debug, Clone, Default)]
162pub struct RawParams {
163    pub strength: Option<f64>,
164    pub variant: Option<String>,
165    pub preset: Option<String>,
166    pub prefilter: Option<String>,
167    pub channel_mode: Option<String>,
168    pub luma_strength: Option<f64>,
169    pub chroma_strength: Option<f64>,
170    pub luma_lambda_ht: Option<f64>,
171    pub chroma_lambda_ht: Option<f64>,
172    pub luma_mismatch_scale: Option<f64>,
173    pub chroma_mismatch_scale: Option<f64>,
174    pub device: Option<String>,
175    pub accelerators: Option<Vec<String>>,
176    pub search_radius: Option<i64>,
177    pub patch_radius: Option<i64>,
178    pub temporal_radius: Option<i64>,
179    pub sigma: Option<f64>,
180    pub sigma_scale: Option<f64>,
181    pub motion_compensation: Option<bool>,
182    pub lambda_ht: Option<f64>,
183    pub lambda_ht_scale: Option<f64>,
184    pub spatial_radius: Option<i64>,
185    pub refine: Option<i64>,
186}
187
188/// Turns a nonnegative script integer into a `u32`, naming `field` in
189/// the error when it is negative.
190fn nonnegative(value: i64, field: &str) -> Result<u32, anyhow::Error> {
191    u32::try_from(value).map_err(|_| anyhow::anyhow!("{field} must not be negative, got {value}"))
192}
193
194/// Resolves an explicit `channel_mode` string into a [`ChannelIntent`],
195/// rejecting anything the source can't support.
196fn parse_channel_mode(raw: &str, layout: FrameLayout) -> Result<ChannelIntent, anyhow::Error> {
197    let intent = match raw.to_ascii_lowercase().as_str() {
198        "luma" => ChannelIntent::Luma,
199        "chroma" => ChannelIntent::Chroma,
200        "lumachroma" => ChannelIntent::LumaChroma,
201        "yuv" => ChannelIntent::YuvFused,
202        other => {
203            anyhow::bail!("unknown channel_mode '{other}', expected one of luma, chroma, lumachroma, yuv");
204        },
205    };
206
207    intent.validate_for_source(layout)?;
208    Ok(intent)
209}
210
211/// Rejects a parameter set on an algorithm that never reads it.
212///
213/// `strength`, its per-plane overrides, `patch_radius`, `search_radius`,
214/// and `variant` only feed the NLM weighting pass, which `NlmTuning`
215/// belongs to. Core's own `DenoiserOptions::to_nlm_params` builds the
216/// `Nl4d` arm from `NlmParams::default()` for exactly this group of
217/// fields, so none of them reach an `Nl4d` run at all, no matter what a
218/// caller sets. `lambda_ht` and `mismatch_scale` and their per-plane
219/// overrides only feed nl4d's temporal grouping stage. Setting one on
220/// the algorithm that ignores it would silently do nothing, which a
221/// script parameter dictionary has no way to warn about on its own, so
222/// this rejects it instead.
223///
224/// `sigma` and `sigma_scale` both pin or nudge the noise level an HQ
225/// front end would otherwise measure. nl4d always runs that front end,
226/// so both are always valid there. Plain `nlmeans` (`variant="fast"`)
227/// has no noise estimator at all, so both are rejected only in that one
228/// case, checked separately below since it depends on `variant` rather
229/// than `algorithm_kind` alone.
230///
231/// `sigma_scale` alongside `sigma` is not rejected here, even though it
232/// does nothing in that combination: the noise estimator that
233/// `sigma_scale` would nudge never runs once `sigma` pins the level.
234/// This mirrors the CLI, which only warns about that combination rather
235/// than erroring (see `--hq-sigma-scale` and `--sigma-scale`), because
236/// `sigma_scale` is a parameter every configuration understands, it
237/// just has nothing left to scale.
238///
239/// `temporal_radius` is not here because it genuinely affects every
240/// algorithm: it sets `PlaneOptions::mode`, which every algorithm
241/// reads. `preset` is not here for the same reason: both algorithms
242/// resolve dials from it.
243///
244/// This runs before the `RUST_MIN_STACK` stack-safety check in
245/// [`plane_options_from`], so `search_radius` on an `Nl4d` call fails
246/// here first rather than tripping that check, which nl4d can never
247/// actually need since it never applies a caller's `search_radius` in
248/// the first place.
249fn reject_mismatched_params(
250    raw: &RawParams,
251    algorithm_kind: AlgorithmKind,
252    variant: NlmeansVariant,
253) -> Result<(), anyhow::Error> {
254    let nlm_only_params: &[(&str, bool)] = &[
255        ("strength", raw.strength.is_some()),
256        ("luma_strength", raw.luma_strength.is_some()),
257        ("chroma_strength", raw.chroma_strength.is_some()),
258        ("patch_radius", raw.patch_radius.is_some()),
259        ("search_radius", raw.search_radius.is_some()),
260        ("variant", raw.variant.is_some()),
261        ("prefilter", raw.prefilter.is_some()),
262        ("motion_compensation", raw.motion_compensation.is_some()),
263    ];
264    let nl4d_only_params: &[(&str, bool)] = &[
265        ("luma_lambda_ht", raw.luma_lambda_ht.is_some()),
266        ("chroma_lambda_ht", raw.chroma_lambda_ht.is_some()),
267        ("luma_mismatch_scale", raw.luma_mismatch_scale.is_some()),
268        ("chroma_mismatch_scale", raw.chroma_mismatch_scale.is_some()),
269        ("lambda_ht", raw.lambda_ht.is_some()),
270        ("lambda_ht_scale", raw.lambda_ht_scale.is_some()),
271        ("spatial_radius", raw.spatial_radius.is_some()),
272        ("refine", raw.refine.is_some()),
273    ];
274
275    match algorithm_kind {
276        AlgorithmKind::Nl4d => {
277            for (name, is_set) in nlm_only_params {
278                if *is_set {
279                    anyhow::bail!(
280                        "{name} has no effect on nl4d, which has no NLM weighting pass to configure"
281                    );
282                }
283            }
284        },
285        AlgorithmKind::Nlmeans => {
286            for (name, is_set) in nl4d_only_params {
287                if *is_set {
288                    anyhow::bail!("{name} has no effect on nlmeans, which only nl4d reads");
289                }
290            }
291            if variant == NlmeansVariant::Fast {
292                if raw.sigma.is_some() {
293                    anyhow::bail!(
294                        "sigma has no effect on nlmeans variant=\"{}\", which has no noise measurement to pin. Set variant=\"hq\" to use sigma",
295                        variant_name(variant)
296                    );
297                }
298                if raw.sigma_scale.is_some() {
299                    anyhow::bail!(
300                        "sigma_scale has no effect on nlmeans variant=\"{}\", which has no noise measurement to nudge. Set variant=\"hq\" to use sigma_scale",
301                        variant_name(variant)
302                    );
303                }
304            }
305        },
306    }
307
308    Ok(())
309}
310
311/// Validates `raw` against `layout` and builds the [`PlaneOptions`] a
312/// [`PlanarDenoiser`](av_denoise_core::PlanarDenoiser) is created from.
313///
314/// Rejects any parameter `algorithm_kind` does not read first, such as
315/// `strength` on nl4d or `lambda_ht` on nlmeans, rather than accepting
316/// and silently ignoring it. See [`reject_mismatched_params`].
317///
318/// Then rejects `search_radius` above 4 when `RUST_MIN_STACK` is unset or
319/// too small, since cubecl's kernel codegen overflows the default 2 MiB
320/// stack and aborts the process at that radius. `filter.rs` raises the
321/// stack before this runs in the real plugin, so this only ever fires
322/// when that step was skipped, and only for nlmeans, since nl4d never
323/// reaches this check with a `search_radius` set at all.
324pub fn plane_options_from(
325    raw: &RawParams,
326    algorithm_kind: AlgorithmKind,
327    layout: FrameLayout,
328) -> Result<PlaneOptions, anyhow::Error> {
329    // Resolved once and read by both algorithms below, exactly like the
330    // CLI's own `--preset`. An explicit `variant`, `temporal_radius`, or
331    // `search_radius` overrides whatever the preset would have picked,
332    // matching `NlmeansArgs::resolve_preset`'s precedence.
333    let preset = match &raw.preset {
334        None => Preset::default(),
335        Some(p) => parse_preset(p)?,
336    };
337
338    // Only `Nlmeans` reads `variant` at all, so `Nl4d` never parses it,
339    // it is rejected as a mismatched parameter below instead if set.
340    let variant = match algorithm_kind {
341        AlgorithmKind::Nlmeans => match raw.variant.as_deref() {
342            None => nlmeans_variant_for(preset),
343            Some(v) => parse_variant(v)?,
344        },
345        AlgorithmKind::Nl4d => NlmeansVariant::Hq,
346    };
347
348    reject_mismatched_params(raw, algorithm_kind, variant)?;
349
350    if let Some(radius) = raw.search_radius
351        && radius > 4
352        && !av_denoise_core::codegen_stack_is_sufficient()
353    {
354        anyhow::bail!(
355            "search_radius {radius} needs a raised stack, but RUST_MIN_STACK is not set. Values above 4 overflow the default 2 MiB stack during kernel codegen"
356        );
357    }
358
359    let intent = match raw.channel_mode.as_deref() {
360        None => ChannelIntent::LumaChroma,
361        Some(mode) => parse_channel_mode(mode, layout)?,
362    };
363
364    let device = match &raw.device {
365        None => Device::default(),
366        Some(s) => s
367            .parse()
368            .map_err(|e| anyhow::anyhow!("invalid device '{s}': {e}"))?,
369    };
370
371    let accelerators = match &raw.accelerators {
372        None => get_default_accelerators(),
373        Some(names) => names
374            .iter()
375            .map(|s| {
376                s.parse::<Accelerator>()
377                    .map_err(|e| anyhow::anyhow!("invalid accelerator '{s}': {e}"))
378            })
379            .collect::<Result<Vec<_>, _>>()?,
380    };
381
382    // Both algorithms resolve their preset-driven temporal radius the
383    // same way the CLI does: an explicit `temporal_radius` overrides
384    // whatever the preset picks. nl4d groups patches across a temporal
385    // window and has no spatial-only mode, but no preset ever resolves
386    // it to 0, so the `radius == 0` arm below only actually triggers
387    // for `nlmeans`, whose `veryfast` preset does.
388    let preset_temporal_radius = match algorithm_kind {
389        AlgorithmKind::Nlmeans => nlmeans_temporal_radius_for(preset),
390        AlgorithmKind::Nl4d => nl4d_temporal_radius_for(preset),
391    };
392
393    let mode = match raw.temporal_radius {
394        None if preset_temporal_radius == 0 => DenoisingMode::Spacial,
395        None => DenoisingMode::Temporal {
396            radius: preset_temporal_radius,
397        },
398        Some(0) => DenoisingMode::Spacial,
399        Some(radius) => DenoisingMode::Temporal {
400            radius: nonnegative(radius, "temporal_radius")?,
401        },
402    };
403
404    let algorithm = match algorithm_kind {
405        AlgorithmKind::Nlmeans => {
406            let tuning = NlmTuning {
407                search_radius: Some(match raw.search_radius {
408                    None => nlmeans_search_radius_for(preset),
409                    Some(r) => nonnegative(r, "search_radius")?,
410                }),
411                patch_radius: raw
412                    .patch_radius
413                    .map(|r| nonnegative(r, "patch_radius"))
414                    .transpose()?,
415                strength: raw.strength.map(|v| v as f32),
416                ..NlmTuning::default()
417            };
418
419            let motion_compensation = match raw.motion_compensation {
420                Some(true) => MotionCompensationMode::from(MotionSearch::default()),
421                Some(false) | None => MotionCompensationMode::None,
422            };
423
424            let prefilter = match &raw.prefilter {
425                None => PrefilterMode::None,
426                Some(s) => {
427                    let mode = parse_prefilter(s)?;
428                    // `parse_prefilter`'s string grammar has no form
429                    // that produces `External`, but the check stays
430                    // here as a boundary guard rather than trusting
431                    // that invariant silently: `External` needs a
432                    // reference frame supplied through
433                    // `push_frame_with_reference`, which this plugin
434                    // has no way to call.
435                    if matches!(mode, PrefilterMode::External) {
436                        anyhow::bail!(
437                            "prefilter 'external' is not supported by av-denoise-vs, which has no way to supply a reference frame"
438                        );
439                    }
440                    mode
441                },
442            };
443
444            match variant {
445                NlmeansVariant::Fast => Algorithm::Nlmeans(NlmeansOptions {
446                    prefilter,
447                    motion_compensation,
448                    tuning,
449                }),
450                NlmeansVariant::Hq => Algorithm::NlmeansHq(NlmeansHqOptions {
451                    nlm: NlmeansOptions {
452                        prefilter,
453                        motion_compensation,
454                        tuning,
455                    },
456                    hq: HqParams {
457                        sigma_override: raw.sigma.map(|v| v as f32),
458                        sigma_scale: raw
459                            .sigma_scale
460                            .map(|v| v as f32)
461                            .unwrap_or_else(|| HqParams::default().sigma_scale),
462                        // A VapourSynth filter has to return the same
463                        // pixels for a frame no matter what order
464                        // frames were requested in, and history-
465                        // dependent estimation breaks that guarantee
466                        // under random access. See `Nl4dOptions`'s own
467                        // `windowed_noise_estimation` field for the
468                        // same reasoning applied to nl4d.
469                        windowed_noise_estimation: true,
470                        ..HqParams::default()
471                    },
472                }),
473            }
474        },
475        AlgorithmKind::Nl4d => Algorithm::Nl4d(Nl4dOptions {
476            // A VapourSynth filter has to return the same pixels for a
477            // frame no matter what order frames were requested in.
478            // window-local estimation computes sigma from only the
479            // frames in the current window, so the fast path and a
480            // `reseed` after random access agree by construction. There
481            // is no reason to expose the stream-history-dependent
482            // temporal EMA here at all.
483            windowed_noise_estimation: true,
484            sigma: raw.sigma.map(|v| v as f32),
485            sigma_scale: raw
486                .sigma_scale
487                .map(|v| v as f32)
488                .unwrap_or_else(|| Nl4dOptions::default().sigma_scale),
489            lambda_ht: raw.lambda_ht.map(|v| v as f32),
490            lambda_ht_scale: raw
491                .lambda_ht_scale
492                .map(|v| v as f32)
493                .unwrap_or_else(|| Nl4dOptions::default().lambda_ht_scale),
494            spatial_radius: match raw.spatial_radius {
495                Some(r) => nonnegative(r, "spatial_radius")?,
496                None => nl4d_spatial_radius_for(preset),
497            },
498            refine: match raw.refine {
499                Some(r) => nonnegative(r, "refine")?,
500                None => Nl4dOptions::default().refine,
501            },
502            ..Nl4dOptions::default()
503        }),
504    };
505
506    Ok(PlaneOptions {
507        accelerators,
508        device,
509        intent,
510        mode,
511        algorithm,
512        luma_strength: raw.luma_strength.map(|v| v as f32),
513        chroma_strength: raw.chroma_strength.map(|v| v as f32),
514        luma_lambda_ht: raw.luma_lambda_ht.map(|v| v as f32),
515        chroma_lambda_ht: raw.chroma_lambda_ht.map(|v| v as f32),
516        luma_mismatch_scale: raw.luma_mismatch_scale.map(|v| v as f32),
517        chroma_mismatch_scale: raw.chroma_mismatch_scale.map(|v| v as f32),
518    })
519}