diff --git a/Cargo.toml b/Cargo.toml index d7393d3..728be6e 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -20,3 +20,8 @@ strum_macros = "0.28" bon = "3" cubecl = { version = "0.10", features = ["std"] } av-denoise-core = { path = "av-denoise-core", version = "0.4.0-alpha7", default-features = false } + +[profile.test] +opt-level = 3 +debug-assertions = true +overflow-checks = true diff --git a/av-denoise-core/Cargo.toml b/av-denoise-core/Cargo.toml index e116dec..cc3129f 100644 --- a/av-denoise-core/Cargo.toml +++ b/av-denoise-core/Cargo.toml @@ -71,3 +71,7 @@ harness = false [[bench]] name = "reseed" harness = false + +[[bench]] +name = "mc_accuracy" +harness = false diff --git a/av-denoise-core/benches/bench_kernels.rs b/av-denoise-core/benches/bench_kernels.rs index a48ddc6..91d6e27 100644 --- a/av-denoise-core/benches/bench_kernels.rs +++ b/av-denoise-core/benches/bench_kernels.rs @@ -29,6 +29,7 @@ use kernels::mc_chain_compose::ChainComposeBench; use kernels::mc_confidence::McConfidenceBench; use kernels::mc_downscale::DownscaleBench; use kernels::mc_warp::WarpBench; +use kernels::mv_regularise::MvRegulariseBench; use kernels::noise_partial::NoisePartialBench; use kernels::pack_wire::PackWireBench; use kernels::temporal_noise_stats::TemporalNoiseStatsBench; @@ -220,11 +221,23 @@ fn run_all(backend: &str, device: &R::Device) { run(ChainComposeBench { client: client.clone(), }); + run(MvRegulariseBench { + client: client.clone(), + }); + for &(ch, ch_name) in CHANNELS { + run(CollabFusedBench { + client: client.clone(), + ch, + ch_name, + split_mv: false, + }); + } for &(ch, ch_name) in CHANNELS { run(CollabFusedBench { client: client.clone(), ch, ch_name, + split_mv: true, }); } for &(ch, ch_name) in CHANNELS { diff --git a/av-denoise-core/benches/kernels/collab_fused.rs b/av-denoise-core/benches/kernels/collab_fused.rs index 310bd64..cfd376d 100644 --- a/av-denoise-core/benches/kernels/collab_fused.rs +++ b/av-denoise-core/benches/kernels/collab_fused.rs @@ -14,13 +14,13 @@ use super::nl4d_geometry::{ CONFIDENCE_VARIANCE, K_MAX, LAMBDA_HT, + MISMATCH_SCALE2, N_FRAMES, NEIGHBOUR_SLOTS, RADIUS, REFINE, SIGMA, SPATIAL_RADIUS, - THSAD, conf_stride, mv_stride, }; @@ -39,12 +39,37 @@ use super::{H, W, block_sync, make_padded_frame, shapes_with_ch, stored_channels /// block skips its comparisons entirely, so leaving it always open here /// measures the worst case. A bench that gates freely would report a /// time well under the real one. +/// +/// `split_mv` picks which of the two motion fields the arm runs on, and +/// the two bracket the real cost of the covering-block search. +/// +/// `false` gives a zeroed field. Every block covering a patch then +/// predicts the same position, all four rectangles coincide, three of +/// them are dropped by the duplicate check and no extra pixel +/// comparison runs. That arm measures the duplicate check on its own. +/// +/// `true` gives each block a vector from its own grid parity, spaced +/// eight pixels apart, which is further than the refine window is wide. +/// The four rectangles covering a patch are then disjoint, nothing +/// deduplicates, and the neighbour search scores four times the +/// positions. That arm is the worst case, and a real motion field lands +/// between the two. pub struct CollabFusedBench { pub client: ComputeClient, pub ch: u32, pub ch_name: &'static str, + pub split_mv: bool, } +/// How far apart two neighbouring blocks' vectors sit in the split +/// field, in pixels. +/// +/// `REFINE` is the rectangle's half-width, so two rectangles stay +/// disjoint once their centres are more than `2 * REFINE` apart. Eight +/// clears that with room and keeps every predicted position well inside +/// a 1080p frame. +const SPLIT_MV_SPACING: i32 = 8; + #[derive(Clone)] pub struct CollabFusedInput { pub ring: Handle, @@ -83,7 +108,18 @@ impl Benchmark for CollabFusedBench { let mv_stride = mv_stride(blocks_x, blocks_y, align); let conf_stride = conf_stride(blocks_x, blocks_y, align); - let mv_data = vec![0i32; (2 * RADIUS * mv_stride) as usize]; + let mut mv_data = vec![0i32; (2 * RADIUS * mv_stride) as usize]; + if self.split_mv { + for t in 0..2 * RADIUS { + for by in 0..blocks_y { + for bx in 0..blocks_x { + let base = (t * mv_stride + (by * blocks_x + bx) * 2) as usize; + mv_data[base] = (bx % 2) as i32 * SPLIT_MV_SPACING; + mv_data[base + 1] = (by % 2) as i32 * SPLIT_MV_SPACING; + } + } + } + } let mv_field = self.client.create_from_slice(i32::as_bytes(&mv_data)); let conf_data = vec![1.0f32; (2 * RADIUS * conf_stride) as usize]; let confidence = self.client.create_from_slice(f32::as_bytes(&conf_data)); @@ -159,7 +195,7 @@ impl Benchmark for CollabFusedBench { CENTRE_SLOT, 0.0f32, 0.0f32, - THSAD, + MISMATCH_SCALE2, LAMBDA_HT, weight_scale(SIGMA, &dct_noise_profile(0.0)), cross_frame_accum_scale(SPATIAL_RADIUS, RADIUS), @@ -185,7 +221,8 @@ impl Benchmark for CollabFusedBench { } fn name(&self) -> String { - format!("collab_fused_1080p_{}", self.ch_name) + let field = if self.split_mv { "_split_mv" } else { "" }; + format!("collab_fused_1080p_{}{field}", self.ch_name) } fn sync(&self) { diff --git a/av-denoise-core/benches/kernels/mod.rs b/av-denoise-core/benches/kernels/mod.rs index dd55e28..00eb543 100644 --- a/av-denoise-core/benches/kernels/mod.rs +++ b/av-denoise-core/benches/kernels/mod.rs @@ -25,6 +25,7 @@ pub mod mc_chain_compose; pub mod mc_confidence; pub mod mc_downscale; pub mod mc_warp; +pub mod mv_regularise; pub mod nl4d_geometry; pub mod noise_partial; pub mod pack_wire; diff --git a/av-denoise-core/benches/kernels/mv_regularise.rs b/av-denoise-core/benches/kernels/mv_regularise.rs new file mode 100644 index 0000000..33c17ce --- /dev/null +++ b/av-denoise-core/benches/kernels/mv_regularise.rs @@ -0,0 +1,92 @@ +use av_denoise_core::nl4d::kernels::nl4d_mv_regularise; +use cubecl::benchmark::Benchmark; +use cubecl::prelude::*; +use cubecl::server::Handle; + +use super::{H, W, block_sync, make_synthetic_frame, shapes_with_ch}; + +const BLKSIZE: u32 = 16; +const STEP: u32 = 8; +const THSAD_PIXEL: f32 = 0.02; +const FIELD_LAMBDA: f32 = 1.0; + +/// The nl4d field regularisation pass over one neighbour at 1080p. +pub struct MvRegulariseBench { + pub client: ComputeClient, +} + +#[derive(Clone)] +pub struct RegulariseInput { + pub centre: Handle, + pub neighbour: Handle, + pub mv_in: Handle, + pub mv_out: Handle, + pub confidence: Handle, +} + +impl Benchmark for MvRegulariseBench { + type Input = RegulariseInput; + type Output = (); + + fn prepare(&self) -> Self::Input { + let blocks = (W.div_ceil(STEP) * H.div_ceil(STEP)) as usize; + let centre = self + .client + .create_from_slice(f32::as_bytes(&make_synthetic_frame(W, H, 1))); + let neighbour = self + .client + .create_from_slice(f32::as_bytes(&make_synthetic_frame(W, H, 1))); + let field: Vec = (0..2 * blocks).map(|i| (i % 7) as i32 - 3).collect(); + let mv_in = self.client.create_from_slice(i32::as_bytes(&field)); + let mv_out = self.client.empty(2 * blocks * size_of::()); + let confidence = self.client.empty(blocks * size_of::()); + RegulariseInput { + centre, + neighbour, + mv_in, + mv_out, + confidence, + } + } + + fn execute(&self, args: Self::Input) -> Result<(), String> { + let blocks_x = W.div_ceil(STEP); + let blocks_y = H.div_ceil(STEP); + let blocks = (blocks_x * blocks_y) as usize; + let block_area = (BLKSIZE * BLKSIZE) as f32; + unsafe { + nl4d_mv_regularise::launch_unchecked::( + &self.client, + CubeCount::new_2d(blocks_x, blocks_y), + CubeDim::new_2d(8, 8), + ArrayArg::from_raw_parts(args.centre.clone(), (W * H) as usize), + ArrayArg::from_raw_parts(args.neighbour.clone(), (W * H) as usize), + ArrayArg::from_raw_parts(args.mv_in.clone(), 2 * blocks), + ArrayArg::from_raw_parts(args.mv_out.clone(), 2 * blocks), + ArrayArg::from_raw_parts(args.confidence.clone(), blocks), + FIELD_LAMBDA * block_area * THSAD_PIXEL, + 0.0, + block_area * THSAD_PIXEL, + W, + H, + BLKSIZE, + STEP, + blocks_x, + blocks_y, + ); + } + Ok(()) + } + + fn name(&self) -> String { + "nl4d_mv_regularise_1080p_luma".to_string() + } + + fn sync(&self) { + block_sync(&self.client); + } + + fn shapes(&self) -> Vec> { + shapes_with_ch(1) + } +} diff --git a/av-denoise-core/benches/kernels/nl4d_geometry.rs b/av-denoise-core/benches/kernels/nl4d_geometry.rs index c86494e..c1ed19e 100644 --- a/av-denoise-core/benches/kernels/nl4d_geometry.rs +++ b/av-denoise-core/benches/kernels/nl4d_geometry.rs @@ -14,7 +14,7 @@ pub const SPATIAL_RADIUS: u32 = 9; /// `collab::MAX_K`, the group size the filter runs at. pub const K_MAX: u32 = 8; /// `Nl4dParams::default().lambda_ht`. -pub const LAMBDA_HT: f32 = 5.3; +pub const LAMBDA_HT: f32 = 5.2; /// `Nl4dParams::default().confidence_variance`, the `use_member_sigma` /// flag `collab_fused` compiles against. pub const CONFIDENCE_VARIANCE: bool = true; @@ -26,10 +26,9 @@ pub const BLK_STEP: u32 = 8; /// (`MotionCompensationMode::Mvtools`'s `blksize`), distinct from /// [`BLK_STEP`] above. pub const BLKSIZE: u32 = 16; -/// `thsad(BLKSIZE, 1.0)` in normalised SAD units (block_area * -/// THSAD_PIXEL, see `crate::nlmeans::motion::thsad`), hand-computed -/// here since that function is crate-private. -pub const THSAD: f32 = (BLKSIZE * BLKSIZE) as f32 * 0.02; +/// `Nl4dParams::default().mismatch_scale` squared, the kernel's +/// `mismatch_scale2` argument. +pub const MISMATCH_SCALE2: f32 = 1.0; /// Frames in the ring a pass reads. pub const N_FRAMES: u32 = 2 * RADIUS + 1; diff --git a/av-denoise-core/benches/mc_accuracy.rs b/av-denoise-core/benches/mc_accuracy.rs new file mode 100644 index 0000000..0e76c5f --- /dev/null +++ b/av-denoise-core/benches/mc_accuracy.rs @@ -0,0 +1,187 @@ +//! Scores nl4d's motion field against synthetic clips with known +//! motion. Prints one table per arm. +//! +//! Run with `cargo bench -p av-denoise-core --bench mc_accuracy -- +//! --device discrete:1 --still brick=/path/to/brick.pgm --still +//! asterisk=/path/to/asterisk.pgm`. With no `--still` it runs on a +//! synthetic texture and says so. + +use std::path::PathBuf; + +use av_denoise_core::nl4d::harness::{Clip, KindScore, MotionClass, Score, Still, score, synthesise}; +use av_denoise_core::nl4d::{Nl4dDenoiser, Nl4dParams}; +use av_denoise_core::nlmeans::{ChannelMode, MotionCompensationMode, NlmParams}; +use cubecl::prelude::*; + +/// Grain levels on the 8-bit scale. +const GRAIN: [f32; 3] = [2.0, 6.0, 12.0]; + +/// A named still. +struct NamedStill { + name: String, + still: Still, +} + +/// One configuration under test. +struct Arm { + name: &'static str, + params: fn() -> Nl4dParams, +} + +/// `Nl4dParams::default` carries `ChannelMode::Yuv`, which expects +/// three interleaved planes per pushed frame. The harness only ever +/// synthesises a single luma plane, so every arm here switches to +/// `ChannelMode::Luma` instead. +fn baseline_params() -> Nl4dParams { + Nl4dParams { + nlm: NlmParams { + channels: ChannelMode::Luma, + ..Nl4dParams::default().nlm + }, + ..Nl4dParams::default() + } +} + +/// `baseline_params` with `field_lambda` overridden, for context against +/// the shipped default. Building from `Nl4dParams::default()` directly +/// would panic with the three-channel default the harness cannot feed. +fn with_lambda_0_5() -> Nl4dParams { + Nl4dParams { + field_lambda: 0.5, + ..baseline_params() + } +} + +/// `baseline_params` with the motion pyramid deepened to three levels, +/// to test whether the extra level earns its added kernel launch. +fn with_pyramid_3() -> Nl4dParams { + let mut p = baseline_params(); + if let MotionCompensationMode::Mvtools { pyramid_levels, .. } = &mut p.nlm.motion_compensation { + *pyramid_levels = 3; + } + p +} + +fn arms() -> Vec { + vec![ + Arm { + name: "baseline", + params: baseline_params, + }, + Arm { + name: "lambda_0.5", + params: with_lambda_0_5, + }, + Arm { + name: "pyramid_3", + params: with_pyramid_3, + }, + ] +} + +fn parse_still(spec: &str) -> Result { + let (name, path) = spec + .split_once('=') + .ok_or_else(|| format!("--still expects name=path, got {spec}"))?; + let bytes = std::fs::read(PathBuf::from(path)).map_err(|e| format!("{path}: {e}"))?; + Ok(NamedStill { + name: name.to_string(), + still: Still::from_pgm(&bytes)?, + }) +} + +fn run_clip(client: &ComputeClient, params: Nl4dParams, clip: &Clip) -> Score { + let refine = params.refine; + let mut d = Nl4dDenoiser::::new(client, params, clip.width, clip.height).expect("construction failed"); + for frame in &clip.frames { + d.push_frame(frame); + let _ = d.denoise_submit().expect("denoise_submit failed"); + } + let snap = d.motion_snapshot().expect("a pass ran once the window filled"); + score(clip, &snap, refine) +} + +fn print_kind(label: &str, k: &KindScore) { + if k.patches == 0 { + return; + } + println!( + " {label:<9} {:>6} corner {:>5.1}% covering {:>5.1}% epe {:>5.2} / p95 {:>5.2} conf {:>4.2}", + k.patches, + 100.0 * k.in_window_rate_corner(), + 100.0 * k.in_window_rate_covering(), + k.epe_mean(), + k.epe_p95(), + k.confidence_median(), + ); +} + +fn run_all(device: &R::Device, stills: &[NamedStill]) { + let client = R::client(device); + for arm in arms() { + println!(); + println!("=== arm: {} ===", arm.name); + for still in stills { + for class in MotionClass::ALL { + for grain in GRAIN { + let params = (arm.params)(); + let clip = synthesise(&still.still, class, params.temporal_radius, grain / 255.0, 7); + let s = run_clip::(&client, params, &clip); + println!(" {:<10} {:<9} grain {grain:>4.0}", still.name, class.label()); + print_kind("plain", &s.plain); + print_kind("boundary", &s.boundary); + print_kind("occluded", &s.occluded); + } + } + } + } +} + +#[derive(clap::Parser, Debug)] +#[command(about = "Motion-field accuracy against synthetic known-motion clips", long_about = None)] +struct Cli { + /// GPU device to bind to. Format: `default`, `discrete[:N]`, + /// `integrated[:N]`, `virtual[:N]`, or `cpu`. + #[arg(long, default_value = "default")] + device: av_denoise_core::Device, + + /// A still to build clips from, as `name=path.pgm`. Repeatable. + #[arg(long = "still")] + stills: Vec, + + /// Swallowed. Cargo passes this when invoking the bench binary. + #[arg(long, hide = true)] + bench: bool, +} + +fn main() { + use clap::Parser; + let cli = Cli::parse(); + + let stills: Vec = if cli.stills.is_empty() { + println!("no --still given, running on a synthetic 256x256 texture"); + vec![NamedStill { + name: "synthetic".to_string(), + still: Still::synthetic(256, 256), + }] + } else { + cli.stills + .iter() + .map(|s| parse_still(s).unwrap_or_else(|e| panic!("{e}"))) + .collect() + }; + + #[cfg(feature = "vulkan")] + { + let device = cli.device.to_wgpu().expect("wgpu device conversion failed"); + println!("device: {device:?}"); + run_all::(&device, &stills); + } + + #[cfg(not(feature = "vulkan"))] + { + let _ = stills; + eprintln!("No GPU backend enabled. Run with --features vulkan"); + std::process::exit(1); + } +} diff --git a/av-denoise-core/benches/motion.rs b/av-denoise-core/benches/motion.rs index 26fbd29..0b04d74 100644 --- a/av-denoise-core/benches/motion.rs +++ b/av-denoise-core/benches/motion.rs @@ -25,6 +25,7 @@ use kernels::mc_block_match_fine::BlockMatchFineBench; use kernels::mc_confidence::McConfidenceBench; use kernels::mc_downscale::DownscaleBench; use kernels::mc_warp::WarpBench; +use kernels::mv_regularise::MvRegulariseBench; use kernels::{CHANNELS, print_header, run}; const W: u32 = 1920; @@ -229,6 +230,9 @@ fn run_kernels(backend: &str, client: &ComputeClient) { run(McConfidenceBench { client: client.clone(), }); + run(MvRegulariseBench { + client: client.clone(), + }); for &(ch, ch_name) in CHANNELS { run(WarpBench { client: client.clone(), diff --git a/av-denoise-core/benches/nl4d_ablation.rs b/av-denoise-core/benches/nl4d_ablation.rs index b50ef9e..bca4aae 100644 --- a/av-denoise-core/benches/nl4d_ablation.rs +++ b/av-denoise-core/benches/nl4d_ablation.rs @@ -54,12 +54,15 @@ const SPATIAL_RADIUS: u32 = 9; const K_MAX: u32 = 8; const BLK_STEP: u32 = 8; const BLKSIZE: u32 = 16; -const THSAD: f32 = (BLKSIZE * BLKSIZE) as f32 * 0.02; +/// `Nl4dParams::default().mismatch_scale` squared, the kernel's +/// `mismatch_scale2` argument. +const MISMATCH_SCALE2: f32 = 1.0; const N_FRAMES: u32 = 2 * RADIUS + 1; const CENTRE_SLOT: u32 = RADIUS; const NEIGHBOUR_SLOTS: [u32; 4] = [0, 1, 3, 4]; const SIGMA: f32 = 0.02; -const LAMBDA_HT: f32 = 5.3; +/// `Nl4dParams::default().lambda_ht`. +const LAMBDA_HT: f32 = 5.2; fn frame_data(g: Geom) -> Vec { let mut data = Vec::with_capacity((g.w * g.h * g.stored) as usize); @@ -210,7 +213,7 @@ impl Rig { CENTRE_SLOT, 0.0f32, 0.0f32, - THSAD, + MISMATCH_SCALE2, LAMBDA_HT, weight_scale(SIGMA, &dct_noise_profile(0.0)), cross_frame_accum_scale(SPATIAL_RADIUS, RADIUS), diff --git a/av-denoise-core/src/collab/kernels/fused.rs b/av-denoise-core/src/collab/kernels/fused.rs index 5c340ae..9415d63 100644 --- a/av-denoise-core/src/collab/kernels/fused.rs +++ b/av-denoise-core/src/collab/kernels/fused.rs @@ -49,53 +49,15 @@ const _: () = assert!( // all of those as compile-time-only, and the shift-insert needs genuine // mutable runtime variables. -/// The extra per-member variance a temporal candidate's motion-block -/// confidence implies, which [`collab_fused`] folds into that member's -/// own noise variance before the threshold reads it. -/// -/// A poorly matched motion block is treated as a noisier observation of -/// the true patch rather than a different patch, so its confidence `c` -/// turns into extra variance instead of an admission decision. -/// -/// `mismatch_thsad` is the SAD threshold that block's confidence score -/// was derived from (see [`crate::nlmeans::motion::thsad`]), multiplied -/// by the caller's `mismatch_scale`, in normalised SAD units. -/// `blksize_area` is the motion block's area in pixels. -/// -/// ```text -/// E^2 = mismatch_thsad^2 * (1 - c) / (1 + c) -/// eps = E / blksize_area -/// sigma_m2 = (pi / 2) * eps^2 -/// ``` -/// -/// The scale is folded into the threshold on the host rather than -/// carried separately, because the two only ever appear multiplied -/// together. It is a scale on the mismatch model alone, not on the -/// confidence score, which stays derived from the unscaled threshold. -/// -/// `c = 1`, a perfect match, gives `sigma_m2 = 0` exactly. Lower -/// confidence inflates it. -/// -/// This never runs for a centre-frame member. Those are not -/// motion-predicted, so there is no mismatch to model, and -/// [`collab_fused`] takes that branch before calling this. -#[cube] -pub(crate) fn mismatch_sigma2(confidence: f32, mismatch_thsad: f32, blksize_area: f32) -> f32 { - let ratio = (1.0f32 - confidence) / (1.0f32 + confidence); - let e2 = mismatch_thsad * mismatch_thsad * ratio; - let eps = f32::sqrt(e2) / blksize_area; - std::f32::consts::FRAC_PI_2 * eps * eps -} - /// The most extra variance a temporal member's mismatch may carry, /// as a multiple of the channel's own variance. /// -/// [`mismatch_sigma2`] derives its result from `mismatch_thsad` and a confidence -/// score, neither of which has any relation to the channel sigma the -/// group weight is normalised against. Left uncapped it makes the -/// retained variance sum, and so the weight, unbounded below, and a -/// weight small enough to round away in the accumulators takes its -/// pixel's only information with it. See +/// A member's extra variance is its own match distance, which on a +/// badly matched patch has no relation to the channel sigma the group +/// weight is normalised against. Left uncapped it makes the retained +/// variance sum, and so the weight, unbounded below, and a weight small +/// enough to round away in the accumulators takes its pixel's only +/// information with it. See /// [`crate::collab::kernels::aggregate::weight_scale`]. /// /// Capping restores the bound. A member here is already a 64 times @@ -105,6 +67,30 @@ pub(crate) fn mismatch_sigma2(confidence: f32, mismatch_thsad: f32, blksize_area /// weight that always survives the conversion to fixed point. pub(crate) const MEMBER_SIGMA2_CAP: f32 = 64.0; +/// The lowest block index whose span contains the patch at `p` on one axis. +/// +/// Block `b` spans `b * step..b * step + blksize`, so the patch +/// `p..p + PATCH_SIZE` needs `b * step + blksize >= p + PATCH_SIZE`. +/// The highest such block is `p / step`, which the caller clamps to the +/// grid and uses as the low end's ceiling. +/// +/// This mirrors `covering_blocks` in the `mc_accuracy` bench's harness +/// module (`av-denoise-core/benches/harness/score.rs`), which the tests +/// below reproduce on the host to check the two stay in step. +#[cube] +fn covering_lo(p: u32, #[comptime] blksize: u32, #[comptime] step: u32) -> u32 { + let past = u32::max(p + PATCH_SIZE, blksize) - blksize; + past.div_ceil(step) +} + +/// The host mirror of [`covering_lo`], for tests that cannot launch a +/// kernel. +#[cfg(test)] +fn covering_lo_host(p: u32, blksize: u32, step: u32) -> u32 { + let past = u32::max(p + PATCH_SIZE, blksize) - blksize; + past.div_ceil(step) +} + /// Groups each reference patch with the patches most similar to it, /// filters the whole group jointly with a hard threshold in the /// transform domain, and scatters every filtered member back into its @@ -154,14 +140,24 @@ pub(crate) const MEMBER_SIGMA2_CAP: f32 = 64.0; /// /// The centre frame contributes the `spatial_radius` rectangle around /// the reference patch, clipped to the frame. Each neighbour -/// contributes the `refine` rectangle around the position the motion -/// field predicts the reference patch moved to, clipped the same way. -/// -/// Clipping the rectangle once is what keeps every candidate a distinct -/// position. Clamping each offset in turn would land several offsets on -/// the same edge position, and admitting a position twice would let one -/// physical patch count as two and look like stronger agreement than -/// the group has. +/// contributes one `refine` rectangle per motion block whose span +/// contains the reference patch, each around the position that block's +/// vector predicts the patch moved to, clipped the same way. A block +/// grid at a step below `blksize` gives several such blocks, and taking +/// all of them means a patch is searched wherever any block covering it +/// points rather than only where its corner block points. +/// +/// Rectangles from different blocks of one neighbour overlap when their +/// vectors are close. A position reached by more than one of them is +/// scored once, by the first rectangle that reaches it, and the later +/// rectangles skip it. +/// +/// Clipping the rectangle once is what keeps every candidate within it a +/// distinct position. Clamping each offset in turn would land several +/// offsets on the same edge position, and admitting a position twice +/// would let one physical patch count as two and look like stronger +/// agreement than the group has. The overlap check across rectangles is +/// the same property held across the blocks of one neighbour. /// /// # Distance /// @@ -169,19 +165,22 @@ pub(crate) const MEMBER_SIGMA2_CAP: f32 = 64.0; /// differences over the whole patch, minus `noise_floor`. `noise_floor` /// is the distance two noisy copies of the same content show by chance, /// so a genuine match is not penalised for the noise it carries. The -/// result is not clamped at zero, because subtracting a constant from -/// every candidate shifts them all equally and leaves the ranking -/// unchanged. +/// result is not clamped at zero for ranking, because subtracting a +/// constant from every candidate shifts them all equally. It is clamped +/// at zero where it becomes a member's mismatch variance below, so +/// `noise_floor` has to be the real expected distance of two noisy +/// copies, `channel_scale * 2 * PATCH_AREA * sum(sigma_c^2)`. /// /// # No admission gate /// /// Every candidate stays in the running whatever its distance, so a /// group fills to `k_max` wherever the search space is that large. -/// `c_min` is a compute saving rather than an admission threshold. A -/// neighbour whose block confidence sits below it never runs the pixel -/// comparison, and its whole rectangle is skipped. The confidence comes -/// from one motion block that every lane of the group shares, so the -/// skip is uniform across the group. +/// `c_min` is a compute saving rather than an admission threshold. The +/// skip is per block. A covering block whose confidence sits below +/// `c_min` never runs the pixel comparison, and its whole rectangle is +/// skipped, while the neighbour's other covering blocks still search. +/// The confidence comes from a motion block that every lane of the +/// group shares, so the skip is uniform across the group. /// /// # Selection /// @@ -225,8 +224,11 @@ pub(crate) const MEMBER_SIGMA2_CAP: f32 = 64.0; /// members carry lands in the higher ones. A coefficient survives a hard /// threshold when its magnitude reaches `lambda_ht` standard deviations /// of its own propagated noise, with [`variance_reg_level`] propagating -/// the per-member variance to each stack level. Both transforms then -/// invert. +/// the per-member variance to each stack level. A member matched in a +/// neighbour frame carries its own match distance as extra variance, +/// `mismatch_scale2 * max(distance, 0) / (3 * PATCH_AREA)`, which is +/// the per-channel, per-pixel mean square of its mismatch, so a poorer +/// match is a noisier observation. Both transforms then invert. /// /// The spatial pass runs as a column DCT in registers, a transpose, and /// a row DCT in registers, because a lane owns a column and the row pass @@ -312,7 +314,7 @@ pub fn collab_fused( centre_slot: u32, noise_floor: f32, c_min: f32, - mismatch_thsad: f32, + mismatch_scale2: f32, lambda_ht: f32, weight_scale: f32, accum_scale: f32, @@ -381,16 +383,21 @@ pub fn collab_fused( // difference. let scale = channel_scale(channels); - // The block a temporal candidate reads its motion vector and - // confidence from depends only on `rx` and `ry`, which are the same - // for every candidate this group scores, so it is worked out once. - let bx = (rx / blk_step).min(blocks_x - 1); - let by = (ry / blk_step).min(blocks_y - 1); - let block = by * blocks_x + bx; - - // The size of the search space, which fixes the group size below. - // Every rectangle contributes distinct positions, and rectangles in - // different frames cannot collide, so this is a plain sum. + // The blocks a temporal candidate reads its motion vectors and + // confidences from depend only on `rx` and `ry`, which are the same + // for every candidate this group scores, so the range is worked out + // once. The corner block, the one the patch's own top-left pixel + // sits in, is `(bx_hi, by_hi)`, and a range whose low end equals its + // high end searches that block alone. + let bx_hi = (rx / blk_step).min(blocks_x - 1); + let by_hi = (ry / blk_step).min(blocks_y - 1); + let bx_lo = u32::min(covering_lo(rx, blksize, blk_step), bx_hi); + let by_lo = u32::min(covering_lo(ry, blksize, blk_step), by_hi); + + // The number of positions actually scored, which fixes the group + // size below. Rectangles in different frames cannot collide, and + // within a frame a repeated position is counted once, so every + // increment is a distinct position. let mut n_live = 0u32; // The spatial rectangle, clipped once. @@ -437,59 +444,110 @@ pub fn collab_fused( cy += 1u32; } - // One clipped rectangle per neighbour, around its motion-predicted - // centre. + // One clipped rectangle per covering block per neighbour, around + // that block's motion-predicted centre. let n_neighbours = comptime!(2 * radius); + // The widest block range `covering_lo` can produce on one axis, so + // the block loops unroll and every `seen_*` index is a constant. + let covers = comptime!(blksize.div_ceil(blk_step)); + let max_rects = comptime!(covers * covers); let mut t = 0u32; while t < n_neighbours { - let conf = confidence[(t * conf_stride + block) as usize]; - // Uniform across the group, because `block` is, so a skipped - // neighbour costs no lane its share of the reduction. No barrier - // sits inside this branch either, so a group that skips a - // neighbour a neighbouring group scores strands nothing. - if conf >= c_min { - let slot = neighbour_slots[t as usize]; - let mv = (t * mv_stride + block * 2u32) as usize; - let px0 = rx as i32 + mv_field[mv]; - let py0 = ry as i32 + mv_field[mv + 1]; - - let t_left = clamp_top_left(px0 - refine as i32, max_x); - let t_right = clamp_top_left(px0 + refine as i32, max_x); - let t_top = clamp_top_left(py0 - refine as i32, max_y); - let t_bot = clamp_top_left(py0 + refine as i32, max_y); - n_live += (t_right - t_left + 1u32) * (t_bot - t_top + 1u32); - - // `t + 1` is the neighbour field's value, one past the - // centre frame's 0. The module-level assert above bounds it - // well inside the six bits `pack_pos_t` gives it. - let packed_t = t + 1u32; - - let mut ny = t_top; - while ny <= t_bot { - let mut nx = t_left; - while nx <= t_right { - let mut partial = 0.0f32; - #[unroll] - for r in 0..PATCH_SIZE { - let px = read_line(ring, nx + sub, ny + r, slot, width, height); - #[unroll] - for c in 0..channels { - let d = current[(r * channels + c) as usize] - px[c as usize]; - partial += d * d; + let slot = neighbour_slots[t as usize]; + // `t + 1` is the neighbour field's value, one past the centre + // frame's 0. The module-level assert above bounds it well inside + // the six bits `pack_pos_t` gives it. + let packed_t = t + 1u32; + + // The rectangles already searched for this neighbour, one slot + // per covering block in visiting order. A slot starts empty, + // `left` above `right`, which no position matches, so a slot + // whose block the scan has not reached yet hides nothing and a + // block the range or `c_min` skips leaves its slot empty. + let mut seen_left = Array::::new(max_rects as usize); + let mut seen_right = Array::::new(max_rects as usize); + let mut seen_top = Array::::new(max_rects as usize); + let mut seen_bot = Array::::new(max_rects as usize); + #[unroll] + for s in 0..max_rects { + seen_left[s as usize] = 1u32; + seen_right[s as usize] = 0u32; + seen_top[s as usize] = 1u32; + seen_bot[s as usize] = 0u32; + } + + #[unroll] + for iy in 0..covers { + #[unroll] + for ix in 0..covers { + let cbx = bx_lo + ix; + let cby = by_lo + iy; + if cbx <= bx_hi && cby <= by_hi { + let block = cby * blocks_x + cbx; + let conf = confidence[(t * conf_stride + block) as usize]; + // Uniform across the group, because `block` is, so a + // skipped block costs no lane its share of the + // reduction. No barrier sits inside this branch + // either, so a group that skips a block a + // neighbouring group scores strands nothing. + if conf >= c_min { + let mv = (t * mv_stride + block * 2u32) as usize; + let px0 = rx as i32 + mv_field[mv]; + let py0 = ry as i32 + mv_field[mv + 1]; + + let t_left = clamp_top_left(px0 - refine as i32, max_x); + let t_right = clamp_top_left(px0 + refine as i32, max_x); + let t_top = clamp_top_left(py0 - refine as i32, max_y); + let t_bot = clamp_top_left(py0 + refine as i32, max_y); + + let mut ny = t_top; + while ny <= t_bot { + let mut nx = t_left; + while nx <= t_right { + let mut covered = false; + #[unroll] + for s in 0..max_rects { + if nx >= seen_left[s as usize] + && nx <= seen_right[s as usize] + && ny >= seen_top[s as usize] + && ny <= seen_bot[s as usize] + { + covered = true; + } + } + if !covered { + n_live += 1u32; + let mut partial = 0.0f32; + #[unroll] + for r in 0..PATCH_SIZE { + let px = read_line(ring, nx + sub, ny + r, slot, width, height); + #[unroll] + for c in 0..channels { + let d = current[(r * channels + c) as usize] - px[c as usize]; + partial += d * d; + } + } + let dist = plane_ssd_reduce8(partial) * scale - noise_floor; + shift_insert8_gated( + &mut best_d, + &mut best_pos, + dist, + pack_pos_t(nx, ny, packed_t), + sub, + base, + ); + } + nx += 1u32; + } + ny += 1u32; } + + seen_left[(iy * covers + ix) as usize] = t_left; + seen_right[(iy * covers + ix) as usize] = t_right; + seen_top[(iy * covers + ix) as usize] = t_top; + seen_bot[(iy * covers + ix) as usize] = t_bot; } - let dist = plane_ssd_reduce8(partial) * scale - noise_floor; - shift_insert8_gated( - &mut best_d, - &mut best_pos, - dist, - pack_pos_t(nx, ny, packed_t), - sub, - base, - ); - nx += 1u32; } - ny += 1u32; } } t += 1u32; @@ -499,7 +557,6 @@ pub fn collab_fused( // member hands every lane every position, once for the whole filter // rather than once per channel. let ref_idx = CUBE_POS_Y * refs_x + ref_x_clamped; - let blksize_area = comptime!(blksize * blksize) as f32; let mut k_use = 1u32; while k_use * 2u32 <= n_live && k_use * 2u32 <= k_max { @@ -538,11 +595,11 @@ pub fn collab_fused( // A centre-frame member is not motion-predicted, so there is // no mismatch to model and it keeps the plain `sigma^2`. if mt > 0u32 { - sig2 = mismatch_sigma2( - confidence[(n * conf_stride + block) as usize], - mismatch_thsad, - blksize_area, - ); + // The member's own distance, floor removed, in the search's + // three-channel-sum units. Per channel and per pixel that + // is the mean square of its mismatch. + let excess = f32::max(plane_shuffle(best_d, base + m), 0.0f32); + sig2 = mismatch_scale2 * excess / comptime!(3 * PATCH_AREA) as f32; } } member_sig2[m as usize] = sig2; @@ -752,3 +809,46 @@ pub fn collab_fused( } } } + +#[cfg(test)] +mod tests { + use super::covering_lo_host; + + /// The host mirror of `covering_blocks` in the `mc_accuracy` bench's + /// harness (`benches/harness/score.rs`). Reproduced here, rather than + /// imported, because that module lives outside the crate as bench-only + /// code and cannot be a test dependency of the library. + /// + /// This pins the kernel's arithmetic against the harness's read of the + /// same geometry rather than launching a real kernel, so it catches the + /// two formulas drifting apart on paper but says nothing about whether + /// [`super::covering_lo`] compiles or runs correctly on a GPU; the + /// integration tests in `nl4d::tests` cover that by driving the whole + /// pipeline. + fn covering_blocks_host(p: u32, blksize: u32, step: u32, blocks: u32) -> (u32, u32) { + let hi = (p / step).min(blocks - 1); + let lo = if p + super::PATCH_SIZE <= blksize { + 0 + } else { + (p + super::PATCH_SIZE - blksize).div_ceil(step) + }; + (lo.min(hi), hi) + } + + #[test] + fn covering_lo_matches_the_harness_across_a_range_of_geometries() { + for (blksize, overlap) in [(16u32, 8u32), (16, 12), (32, 24), (8, 4), (16, 0)] { + let step = blksize - overlap; + let blocks = 8u32; + for p in (0..blocks * step).step_by(3) { + let (expect_lo, hi) = covering_blocks_host(p, blksize, step, blocks); + let got_lo = covering_lo_host(p, blksize, step).min(hi); + assert_eq!( + got_lo, expect_lo, + "blksize={blksize} step={step} p={p}: covering_lo disagrees with the \ + harness's covering_blocks" + ); + } + } + } +} diff --git a/av-denoise-core/src/collab/kernels/transforms.rs b/av-denoise-core/src/collab/kernels/transforms.rs index e6b4ada..6079fa9 100644 --- a/av-denoise-core/src/collab/kernels/transforms.rs +++ b/av-denoise-core/src/collab/kernels/transforms.rs @@ -312,7 +312,7 @@ pub(crate) fn haar_variance_ladder(sig2: &[f32], k_use: u32) -> Vec { out } -#[cfg(test)] +#[cfg(all(test, any(feature = "vulkan", feature = "metal")))] mod tests { use super::*; diff --git a/av-denoise-core/src/collab/tests/fused.rs b/av-denoise-core/src/collab/tests/fused.rs index 991ffe1..39d8016 100644 --- a/av-denoise-core/src/collab/tests/fused.rs +++ b/av-denoise-core/src/collab/tests/fused.rs @@ -29,18 +29,15 @@ const SPATIAL_RADIUS: u32 = 4; /// for. const K_MAX: u32 = 8; -/// Motion-block side length, the value the mismatch variance is scored -/// against. +/// Motion-block side length. The kernel searches every block whose +/// `blksize` span contains a patch, and does not score the mismatch +/// variance against it. const BLKSIZE: u32 = 16; /// Motion-block stride. It stays at `PATCH_SIZE` so a block boundary /// lines up with a patch boundary. const BLK_STEP: u32 = 8; -/// `thsad(BLKSIZE, 1.0)` in normalised SAD units, the same value a real -/// caller gets at this block size and the default scale. -const THSAD: f32 = (BLKSIZE * BLKSIZE) as f32 * 0.02; - /// The noise level the filter is told to shrink against. /// /// Small enough against content in `[0, 1]` that the threshold keeps a @@ -48,7 +45,11 @@ const THSAD: f32 = (BLKSIZE * BLKSIZE) as f32 * 0.02; /// sides of the keep decision are exercised. const SIGMA: f32 = 0.02; +/// A fixed hard-threshold multiplier, pinned independently of /// `Nl4dParams::default().lambda_ht`. +/// +/// Several tests in this file recorded their expected output at this +/// value, so it stays fixed even when the shipped default moves. const LAMBDA_HT: f32 = 5.3; /// [`make_unique_frame`] rescaled into `[0, 1]`. @@ -90,17 +91,16 @@ struct Setup { k_max: u32, sigma: f32, lambda_ht: f32, - /// Whether a temporal member's motion-block confidence inflates its - /// own noise variance. + /// Whether a temporal member's own match distance inflates its + /// noise variance. confidence_variance: bool, /// Residual correlation the noise profile is built for. `0.0` gives /// the all-ones profile most runs use. rho: f32, - /// The SAD threshold a temporal member's mismatch variance is derived - /// from. [`THSAD`] is what a real caller passes at the default block - /// size, and raising it is how a run drives that variance far above - /// the channel sigma. - thsad: f32, + /// The multiplier on a temporal member's mismatch variance, `1.0` at + /// its default. Squared before it reaches the kernel, see + /// [`crate::nl4d::params::Nl4dParams::mismatch_scale`]. + mismatch_scale: f32, /// A profile buffer supplied outright, bypassing /// [`dct_noise_profile`]. The weight scale still follows whatever /// profile is in force. @@ -139,7 +139,7 @@ impl Setup { lambda_ht: LAMBDA_HT, confidence_variance: true, rho: 0.0, - thsad: THSAD, + mismatch_scale: 1.0, profile_override: None, kaiser_beta: 0.0, } @@ -438,7 +438,7 @@ fn run_fused(s: &Setup) -> Aggregated { s.centre_slot, s.noise_floor, s.c_min, - s.thsad, + s.mismatch_scale * s.mismatch_scale, s.lambda_ht, weight_scale(s.sigma, &profile), s.accum_scale(), @@ -740,8 +740,20 @@ fn cross_frame_setup(width: u32, height: u32, radius: u32) -> Setup { } /// The whole temporal path at once: the `c_min` skip, the per-member -/// mismatch variance derived from a packed neighbour index, and the -/// scatter into each member's own region of the accumulator ring. +/// mismatch variance derived from the member's own match distance, and +/// the scatter into each member's own region of the accumulator ring. +/// +/// Recorded with the covering-block search, every block covering a +/// patch contributes a rectangle. `cross_frame_setup` gives every block +/// its own vector, so the search reaches positions the corner block +/// alone never pointed at. +/// +/// Re-recorded for the switch from motion-block confidence to a +/// member's own match distance. The digest below comes from this +/// kernel's own output, not a second implementation, because none +/// exists for the new mechanism. [`assert_matches_recorded`]'s warning +/// about comparing a kernel to itself is about a silently-broken shader +/// producing zeros, and this recording carries real, non-zero coverage. #[test] fn fused_reproduces_recorded_output_across_frames() { let s = cross_frame_setup(64, 64, 2); @@ -749,19 +761,19 @@ fn fused_reproduces_recorded_output_across_frames() { "cross frame", &run_fused(&s), &Digest { - covered: 12928, - pixel_mean: 0.319278212107, - pixel_rms: 0.462380832227, - weight_mean: 1149.191924642, + covered: 15800, + pixel_mean: 0.398917931934, + pixel_rms: 0.518592139022, + weight_mean: 1199.919938422, probes: [ - 0.839722565729, - 0.574316714978, - 0.298724122489, - 0.000000000000, - 0.774649096602, - 0.000000000000, + 0.838003113388, + 0.574141517596, + 0.299827186817, 0.000000000000, + 0.774458945874, 0.000000000000, + 0.236727453142, + 0.979726340630, ], }, ); @@ -772,6 +784,9 @@ fn fused_reproduces_recorded_output_across_frames() { /// `use_member_sigma` is a `#[comptime]` flag, so it compiles a second /// program, and the arm with it off is the one that checks the threshold /// still reads a plain `sigma^2` per member. +/// +/// Recorded with the covering-block search, every block covering a +/// patch contributes a rectangle. #[test] fn fused_reproduces_recorded_output_without_the_mismatch_variance() { let mut s = cross_frame_setup(64, 64, 2); @@ -780,19 +795,19 @@ fn fused_reproduces_recorded_output_without_the_mismatch_variance() { "cross frame, flat sigma", &run_fused(&s), &Digest { - covered: 12928, - pixel_mean: 0.319277061395, - pixel_rms: 0.462378164801, - weight_mean: 1242.592593316, + covered: 15800, + pixel_mean: 0.398918693763, + pixel_rms: 0.518592557620, + weight_mean: 1244.444444987, probes: [ - 0.839714050293, - 0.574348068237, - 0.298727416992, - 0.000000000000, - 0.774412972586, - 0.000000000000, + 0.838030815125, + 0.574148050944, + 0.299845377604, 0.000000000000, + 0.774438040597, 0.000000000000, + 0.236724853516, + 0.979728698730, ], }, ); @@ -934,22 +949,17 @@ fn zero_sigma_hands_every_member_back_unchanged() { } /// A temporal member's mismatch variance has no relation to the channel -/// sigma the group weight is normalised against, so a badly matched group -/// has no lower bound on its weight (see +/// sigma the group weight is normalised against, so a badly matched +/// group has no lower bound on its weight (see /// [`crate::collab::kernels::aggregate::weight_scale`]). Push that /// variance up far enough and the weight stops being representable at /// all, and a group that reaches the accumulators as nothing leaves a /// covered pixel with an empty weight sum, which normalisation can only /// render as black. /// -/// Zero confidence is the worst case `mismatch_sigma2` models, and -/// `thsad` scales the variance it implies. The radii are the shipped -/// defaults, so the run counts in the same fixed point a real cross-frame -/// pass does rather than the finer one a small search would pick. Every -/// pixel a reference covers has to keep carrying weight across all of it. -/// -/// The last two rungs are past anything a caller would ask for, which is -/// the point: they run the mismatch variance so far past +/// Every neighbour here holds content unrelated to the centre, so every +/// temporal member's distance is large, and `mismatch_scale` multiplies +/// the variance it implies. The last rungs run it so far past /// [`crate::collab::kernels::fused::MEMBER_SIGMA2_CAP`] that only the cap /// is holding the weight inside the fixed point at all. #[test] @@ -957,15 +967,11 @@ fn a_badly_matched_group_still_reaches_the_accumulators() { let (w, h) = (32u32, 32u32); let counts = reference_cover_counts(w, h); - for scale in [1.0f32, 64.0, 1024.0, 4096.0] { + for scale in [1.0f32, 8.0, 32.0, 64.0] { let mut s = cross_frame_setup(w, h, 2); s.spatial_radius = 9; - s.confidence.fill(0.0); - // The confidence floor has to come down with it, or the groups - // are skipped before they are ever scored and the run says - // nothing about their weights. s.c_min = 0.0; - s.thsad = THSAD * scale; + s.mismatch_scale = scale; let got = run_fused(&s); let base = s.centre_slot as usize * s.pixels(); @@ -975,7 +981,7 @@ fn a_badly_matched_group_still_reaches_the_accumulators() { } assert!( got.wsum[base + idx] > 0, - "thsad scale {scale}: {count} references cover pixel {idx} and its weight \ + "mismatch scale {scale}: {count} references cover pixel {idx} and its weight \ sum is still {}", got.wsum[base + idx], ); @@ -991,12 +997,11 @@ fn a_windowed_badly_matched_group_still_reaches_the_accumulators() { let (w, h) = (32u32, 32u32); let counts = reference_cover_counts(w, h); - for scale in [1.0f32, 64.0, 1024.0, 4096.0] { + for scale in [1.0f32, 8.0, 32.0, 64.0] { let mut s = cross_frame_setup(w, h, 2); s.spatial_radius = 9; - s.confidence.fill(0.0); s.c_min = 0.0; - s.thsad = THSAD * scale; + s.mismatch_scale = scale; s.kaiser_beta = 2.0; let got = run_fused(&s); @@ -1007,7 +1012,8 @@ fn a_windowed_badly_matched_group_still_reaches_the_accumulators() { } assert!( got.wsum[base + idx] > 0, - "thsad scale {scale}: {count} references cover pixel {idx} and its weight sum is still {} with the window on", + "mismatch scale {scale}: {count} references cover pixel {idx} and its weight \ + sum is still {} with the window on", got.wsum[base + idx], ); } @@ -1456,12 +1462,10 @@ fn higher_rho_retains_more_noise_on_a_flat_field() { /// A centre-frame member never picks up a mismatch variance. /// /// Every neighbour here is gated out by `c_min`, so every member of -/// every group comes from the centre frame, and the confidence buffer -/// holds `0.0`, the value that derives the largest mismatch variance -/// there is. Turning `confidence_variance` on must therefore change -/// nothing at all. A kernel that fed a centre-frame member through -/// [`crate::collab::kernels::fused::mismatch_sigma2`] would inflate -/// every threshold in the frame and move every pixel. +/// every group comes from the centre frame. Turning `confidence_variance` +/// on must therefore change nothing at all. A kernel that computed a +/// mismatch variance for a centre-frame member would inflate every +/// threshold in the frame and move every pixel. #[test] fn centre_frame_members_ignore_the_confidence_field() { let (w, h) = (64u32, 64u32); diff --git a/av-denoise-core/src/denoiser.rs b/av-denoise-core/src/denoiser.rs index 66af79b..a0a3f63 100644 --- a/av-denoise-core/src/denoiser.rs +++ b/av-denoise-core/src/denoiser.rs @@ -258,6 +258,8 @@ pub struct Nl4dOptions { /// estimation breaks that guarantee under random access. See /// [`HqParams::windowed_noise_estimation`]. pub windowed_noise_estimation: bool, + /// See [`crate::nl4d::Nl4dParams::field_lambda`]. + pub field_lambda: f32, } impl Default for Nl4dOptions { @@ -281,6 +283,7 @@ impl Default for Nl4dOptions { confidence_variance: defaults.confidence_variance, kaiser_beta: defaults.kaiser_beta, windowed_noise_estimation: false, + field_lambda: defaults.field_lambda, } } } @@ -311,21 +314,22 @@ impl Nl4dOptions { /// noise and more fine detail with it, so the value is a trade rather /// than an optimum. /// -/// Luma gets 5.3, picked by eye from rendered comparisons on real grain -/// and deliberately biased toward keeping detail. Higher values remove -/// visibly more noise, but not enough to be worth what they cost in -/// texture. +/// Luma and chroma values are picked by eye from a ladder of renders against real +/// film grain, accepting more lost detail in exchange for less remaining +/// noise on heavy grain. Separately confirmed not to over-filter near-clean +/// animation. The reason why we're going a bit heavier on high noise is because +/// the encoders end up reducing that detail _more_ than the denoiser does if +/// that extra entropy remains in and overall produces a worse final image. /// /// `ChannelMode::Yuv` reads the luma value, on the same "a fused pass is /// dominated by luma" assumption [`hq_default_strength`] /// makes for its own Yuv case. /// -/// Chroma gets 4.2, picked the same way from the chroma residuals with -/// luma pinned at 5.3. +/// Luma and the fused Yuv mode use 5.2, and chroma uses 3.4. pub fn nl4d_default_lambda_ht(channels: ChannelMode) -> f32 { match channels { - ChannelMode::Luma | ChannelMode::Yuv => 5.3, - ChannelMode::Chroma => 4.2, + ChannelMode::Luma | ChannelMode::Yuv => 5.2, + ChannelMode::Chroma => 3.4, } } @@ -679,6 +683,7 @@ fn build_engine( mismatch_scale: opts.mismatch_scale, confidence_variance: opts.confidence_variance, kaiser_beta: opts.kaiser_beta, + field_lambda: opts.field_lambda, }; let denoiser = Nl4dDenoiser::with_output_format(client, nl4d_params, width, height, output_format) @@ -1407,8 +1412,8 @@ mod options_tests { let luma = nl4d_default_lambda_ht(ChannelMode::Luma); let chroma = nl4d_default_lambda_ht(ChannelMode::Chroma); - assert!((luma - 5.3).abs() < f32::EPSILON); - assert!((chroma - 4.2).abs() < f32::EPSILON); + assert!((luma - 5.2).abs() < f32::EPSILON); + assert!((chroma - 3.4).abs() < f32::EPSILON); assert!( (chroma - luma).abs() > f32::EPSILON, "the two planes should not resolve to the same default" @@ -1430,8 +1435,8 @@ mod options_tests { let luma = resolve_lambda_ht(&opts, ChannelMode::Luma).expect("the default scale is in range"); let chroma = resolve_lambda_ht(&opts, ChannelMode::Chroma).expect("the default scale is in range"); - assert!((luma - 5.3).abs() < f32::EPSILON, "got {luma}"); - assert!((chroma - 4.2).abs() < f32::EPSILON, "got {chroma}"); + assert!((luma - 5.2).abs() < f32::EPSILON, "got {luma}"); + assert!((chroma - 3.4).abs() < f32::EPSILON, "got {chroma}"); } #[test] diff --git a/av-denoise-core/src/frame/mod.rs b/av-denoise-core/src/frame/mod.rs index 822bbb6..347f026 100644 --- a/av-denoise-core/src/frame/mod.rs +++ b/av-denoise-core/src/frame/mod.rs @@ -1401,12 +1401,12 @@ mod cli_options_tests { // ...but resolving each through the same function construction // uses (`nl4d_default_lambda_ht`, see `src/denoiser.rs`) gives // luma and chroma different values, which is the whole point of - // a caller passing no flags at all getting both calibrated + // a caller passing no flags at all getting both per-plane // defaults. let luma_default = crate::nl4d_default_lambda_ht(ChannelMode::Luma); let chroma_default = crate::nl4d_default_lambda_ht(ChannelMode::Chroma); - assert!((luma_default - 5.3).abs() < f32::EPSILON); - assert!((chroma_default - 4.2).abs() < f32::EPSILON); + assert!((luma_default - 5.2).abs() < f32::EPSILON); + assert!((chroma_default - 3.4).abs() < f32::EPSILON); assert!((chroma_default - luma_default).abs() > f32::EPSILON); } } diff --git a/av-denoise-core/src/frame/tests.rs b/av-denoise-core/src/frame/tests.rs index 6713934..d5bf3da 100644 --- a/av-denoise-core/src/frame/tests.rs +++ b/av-denoise-core/src/frame/tests.rs @@ -348,6 +348,15 @@ mod reseed { .unwrap_or(0) } + /// The count of samples whose absolute difference between two + /// same-sized byte planes exceeds `threshold`. + fn count_exceeding(a: &[u8], b: &[u8], threshold: i32) -> usize { + a.iter() + .zip(b.iter()) + .filter(|&(&x, &y)| (x as i32 - y as i32).abs() > threshold) + .count() + } + /// The nl4d mirror of /// [`reseed_matches_the_streaming_output_at_both_clip_edges`], where /// the widened window clamps at both ends of the clip. @@ -376,19 +385,41 @@ mod reseed { assert_eq!(got.u, streamed[last].u, "u mismatch at the ahead edge"); assert_eq!(got.v, streamed[last].v, "v mismatch at the ahead edge"); - // The bound this leading-edge padding difference stays within: - // full-range luma codes span 255, and the extra duplicated - // history `reseed` folds in at the clip's first frame moves the - // result by at most a handful of 8-bit codes, well short of a - // bound that would let a genuine regression through. - const BEHIND_EDGE_TOLERANCE: i32 = 8; + // Two bounds cover this leading-edge padding difference, because + // it has a known shape rather than an unknown one. `reseed` and + // a fresh stream fold different amounts of duplicated history + // into nl4d's cross-frame accumulator right at the clip's first + // frame, and inside that padded region a hard-threshold + // coefficient can sit close enough to its cutoff that the two + // paths land it on opposite sides. That flips the reconstruction + // of a couple of pixels by their own magnitude while leaving the + // rest of the plane alone. `BEHIND_EDGE_TOLERANCE` is a + // worst-pixel bound, sized well under the full 255-code range + // so a real regression would still trip it. `BEHIND_EDGE_OUTLIER_LIMIT` + // is the original, tighter bound of 8 kept as a count instead of + // a ceiling: at most a handful of samples may cross it, and a + // real regression that moved the bulk of the plane would push + // far more samples past it than that. + // + // Measured against this fixture, the actual worst-pixel diff was + // 10, with exactly 1 sample exceeding 8, so both bounds carry + // headroom over what was observed. + const BEHIND_EDGE_TOLERANCE: i32 = 16; + const BEHIND_EDGE_OUTLIER_THRESHOLD: i32 = 8; + const BEHIND_EDGE_OUTLIER_LIMIT: usize = 4; let mut d = PlanarDenoiser::create(&opts, layout()).unwrap(); let got = d.reseed(&window_of_span(&frames, 0, span)).unwrap(); let luma_diff = max_abs_diff(&got.y, &streamed[0].y); + let luma_outliers = count_exceeding(&got.y, &streamed[0].y, BEHIND_EDGE_OUTLIER_THRESHOLD); assert!( luma_diff <= BEHIND_EDGE_TOLERANCE, "luma at the behind edge (k=0) drifted too far from streaming: max abs diff {luma_diff}" ); + assert!( + luma_outliers <= BEHIND_EDGE_OUTLIER_LIMIT, + "luma at the behind edge (k=0) drifted too far from streaming across too much of the \ + plane: {luma_outliers} samples exceeded {BEHIND_EDGE_OUTLIER_THRESHOLD}" + ); } /// After an nl4d `reseed`, ordinary sequential `push`/`recv` must diff --git a/av-denoise-core/src/nl4d/denoiser.rs b/av-denoise-core/src/nl4d/denoiser.rs index 67c62fb..fdafa0f 100644 --- a/av-denoise-core/src/nl4d/denoiser.rs +++ b/av-denoise-core/src/nl4d/denoiser.rs @@ -2,6 +2,8 @@ use cubecl::prelude::*; use cubecl::server::Handle; use super::params::Nl4dParams; +use super::regularise::run_regularise; +use super::snapshot::{LastFields, MotionSnapshot, read_snapshot}; use crate::collab::geometry::{fused_cubes_x, ref_count, refs_along}; use crate::collab::kernels::aggregate::{ collab_normalise, @@ -12,8 +14,9 @@ use crate::collab::kernels::aggregate::{ }; use crate::collab::kernels::fused::collab_fused; use crate::collab::kernels::transforms::dct_noise_profile; -use crate::collab::{MAX_K, PATCH_SIZE}; +use crate::collab::{MAX_K, PATCH_AREA, PATCH_SIZE}; use crate::denoiser::{DenoiserError, FrameOutput, OutputFormat}; +use crate::nlmeans::kernels::helpers::channel_scale_host; use crate::nlmeans::{ BLOCK_X, BLOCK_Y, @@ -129,6 +132,15 @@ pub struct Nl4dDenoiser { /// [`Self::run_collab_stage`]'s return value rather than checking it /// themselves. passes_run: u32, + /// The field buffers the last pass handed the fused kernel, for + /// [`Self::motion_snapshot`]. + last_fields: Option, + /// See [`Nl4dParams::field_lambda`]. + field_lambda: f32, + /// The regularised motion field and confidence, laid out like the + /// front end's own, allocated only when `field_lambda > 0.0`. + reg_mv: Option, + reg_conf: Option, } impl Nl4dDenoiser { @@ -225,6 +237,19 @@ impl Nl4dDenoiser { }, }; + // `motion_ctx()` panics without motion compensation, and + // `validate` above already requires it, so this is safe here. + let (reg_mv, reg_conf) = if params.field_lambda > 0.0 { + let mc = front.motion_ctx(); + let neighbours = 2 * params.temporal_radius as u64; + ( + Some(client.empty((neighbours * mc.mv_field_bytes_per_neighbour()) as usize)), + Some(client.empty((neighbours * mc.confidence_bytes_per_neighbour()) as usize)), + ) + } else { + (None, None) + }; + Ok(Self { front, width, @@ -251,6 +276,10 @@ impl Nl4dDenoiser { output_format, wire_outputs, passes_run: 0, + last_fields: None, + field_lambda: params.field_lambda, + reg_mv, + reg_conf, }) } @@ -373,6 +402,27 @@ impl Nl4dDenoiser { self.front.reset_stream_state(); self.next_output_slot = 0; self.passes_run = 0; + self.last_fields = None; + } + + /// The motion field and confidence the last pass gave the fused + /// kernel, or `None` before any pass has run. + /// + /// This is a synchronous readback for measurement tooling, not a + /// stable interface. + #[doc(hidden)] + pub fn motion_snapshot(&self) -> Option { + let fields = self.last_fields.as_ref()?; + let mc = self.front.motion_ctx(); + Some(read_snapshot( + self.front.compute_client(), + fields, + self.temporal_radius, + mc.blocks_x, + mc.blocks_y, + mc.step, + mc.blksize, + )) } /// How many tail frames [`Self::flush`] must emit for the stream @@ -498,11 +548,14 @@ impl Nl4dDenoiser { let blksize = mc.blksize; let blocks_x = mc.blocks_x; let blocks_y = mc.blocks_y; - // The kernel takes the two multiplied together, see - // `mismatch_sigma2`. The confidence score itself stays derived - // from the unscaled threshold, which is why this is applied here - // rather than inside the front end. - let mismatch_thsad = self.front.thsad_value() * self.mismatch_scale; + // The variance grows with the square of the scale, see + // `Nl4dParams::mismatch_scale`. + let mismatch_scale2 = self.mismatch_scale * self.mismatch_scale; + // The distance two noisy copies of one patch show by chance, in + // the search's channel-scaled units. A member's mismatch + // variance is its distance past this. + let sigma2_sum: f32 = sigma_host[..channels_count as usize].iter().map(|s| s * s).sum(); + let noise_floor = channel_scale_host(channels_count) * 2.0 * PATCH_AREA as f32 * sigma2_sum; // See the doc comment above for why these two slots are what // this pass clears and completes. `total_frames` is added @@ -515,6 +568,36 @@ impl Nl4dDenoiser { let pass_index = self.passes_run; self.passes_run += 1; + // The field the fused kernel reads, the regularised one when the + // pass is on. + let (mv_field, confidence) = match (self.reg_mv.as_ref(), self.reg_conf.as_ref()) { + (Some(mv), Some(conf)) => { + run_regularise::( + &client, + mc, + view, + self.width, + self.height, + self.field_lambda, + self.front.sad_noise_floor_value(), + self.front.thsad_value(), + mv, + conf, + ) + .map_err(DenoiserError::Other)?; + (mv.clone(), conf.clone()) + }, + _ => (view.mv_field.clone(), view.confidence.clone()), + }; + + self.last_fields = Some(LastFields { + mv_field: mv_field.clone(), + confidence: confidence.clone(), + mv_stride: view.mv_stride, + conf_stride: view.conf_stride, + neighbours, + }); + unsafe { if pass_index == 0 { // Clearing the whole ring in one dispatch would need @@ -562,8 +645,8 @@ impl Nl4dDenoiser { collab_dim, stored_ch as usize, ArrayArg::from_raw_parts(view.input.clone(), ring_len), - ArrayArg::from_raw_parts(view.mv_field.clone(), mv_len.max(1)), - ArrayArg::from_raw_parts(view.confidence.clone(), conf_len.max(1)), + ArrayArg::from_raw_parts(mv_field.clone(), mv_len.max(1)), + ArrayArg::from_raw_parts(confidence.clone(), conf_len.max(1)), ArrayArg::from_raw_parts(neighbour_slots_buf, view.neighbour_slots.len().max(1)), ArrayArg::from_raw_parts(self.sigma_buf.clone(), stored_ch as usize), ArrayArg::from_raw_parts(self.dct_profile_buf.clone(), 8), @@ -572,14 +655,9 @@ impl Nl4dDenoiser { ArrayArg::from_raw_parts(self.wsum.clone(), wsum_ring_len), ArrayArg::from_raw_parts(self.group_weight.clone(), refs), centre_slot, - // `collab_fused` has no admission gate (see its own doc - // comment), so a constant subtracted from every - // candidate's distance can never change which ones the - // selection picks. Any value is exact here; 0.0 is the - // simplest one that says so. - 0.0f32, + noise_floor, self.c_min, - mismatch_thsad, + mismatch_scale2, self.lambda_ht, wnorm, self.accum_scale, diff --git a/av-denoise-core/src/nl4d/harness/mod.rs b/av-denoise-core/src/nl4d/harness/mod.rs new file mode 100644 index 0000000..bf18097 --- /dev/null +++ b/av-denoise-core/src/nl4d/harness/mod.rs @@ -0,0 +1,13 @@ +//! Synthetic clips with known motion, and the scores that compare a +//! motion field against them. +//! +//! This exists for the `mc_accuracy` bench. It is not a stable +//! interface. + +#![doc(hidden)] + +mod score; +mod synth; + +pub use score::{KindScore, Score, score}; +pub use synth::{Clip, MotionClass, Still, synthesise}; diff --git a/av-denoise-core/src/nl4d/harness/score.rs b/av-denoise-core/src/nl4d/harness/score.rs new file mode 100644 index 0000000..84f932d --- /dev/null +++ b/av-denoise-core/src/nl4d/harness/score.rs @@ -0,0 +1,338 @@ +use super::synth::Clip; +use crate::collab::PATCH_SIZE; +use crate::collab::geometry::{ref_pos, refs_along}; +use crate::nl4d::MotionSnapshot; + +/// The inclusive range of blocks whose `b * step..b * step + blksize` +/// span contains the patch `p..p + PATCH_SIZE`, clamped to the grid. +/// +/// When `step == blksize` and the patch straddles a tile boundary, no +/// block fully contains it. The corner block is returned as the best +/// available in that case, because a consumer still needs something to +/// search. +pub fn covering_blocks(p: u32, blksize: u32, step: u32, blocks: u32) -> (u32, u32) { + let hi = (p / step).min(blocks - 1); + let lo = if p + PATCH_SIZE <= blksize { + 0 + } else { + (p + PATCH_SIZE - blksize).div_ceil(step) + }; + (lo.min(hi), hi) +} + +/// How a patch's ground truth classifies it. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum PatchKind { + Plain, + Boundary, + Occluded, +} + +/// Aggregated results for one patch kind. +#[derive(Debug, Clone, Default)] +pub struct KindScore { + pub patches: usize, + pub in_window_corner: usize, + pub in_window_covering: usize, + pub epe: Vec, + pub confidence: Vec, +} + +impl KindScore { + pub fn in_window_rate_corner(&self) -> f64 { + if self.patches == 0 { + 0.0 + } else { + self.in_window_corner as f64 / self.patches as f64 + } + } + + pub fn in_window_rate_covering(&self) -> f64 { + if self.patches == 0 { + 0.0 + } else { + self.in_window_covering as f64 / self.patches as f64 + } + } + + pub fn epe_mean(&self) -> f64 { + if self.epe.is_empty() { + 0.0 + } else { + self.epe.iter().map(|&e| e as f64).sum::() / self.epe.len() as f64 + } + } + + pub fn epe_p95(&self) -> f64 { + percentile(&self.epe, 0.95) + } + + pub fn confidence_median(&self) -> f64 { + percentile(&self.confidence, 0.5) + } +} + +fn percentile(values: &[f32], q: f64) -> f64 { + if values.is_empty() { + return 0.0; + } + let mut sorted = values.to_vec(); + sorted.sort_by(|a, b| a.partial_cmp(b).expect("no NaN in scores")); + let idx = ((sorted.len() - 1) as f64 * q).round() as usize; + sorted[idx] as f64 +} + +/// The full score of one field against one clip. +#[derive(Debug, Clone, Default)] +pub struct Score { + pub plain: KindScore, + pub boundary: KindScore, + pub occluded: KindScore, +} + +impl Score { + fn kind_mut(&mut self, kind: PatchKind) -> &mut KindScore { + match kind { + PatchKind::Plain => &mut self.plain, + PatchKind::Boundary => &mut self.boundary, + PatchKind::Occluded => &mut self.occluded, + } + } +} + +/// The largest-axis distance between a truth and an integer vector. +fn endpoint_error(truth: [f32; 2], v: [i32; 2]) -> f32 { + (truth[0] - v[0] as f32).abs().max((truth[1] - v[1] as f32).abs()) +} + +/// Scores `snap` against `clip` over nl4d's reference grid and every +/// neighbour. +/// +/// A patch is in window when its truth lies within `refine` pixels of +/// the vector on both axes. The corner reading uses the block whose +/// corner the patch sits on, the block nl4d reads today. The covering +/// reading takes the best of every block that covers the patch. +pub fn score(clip: &Clip, snap: &MotionSnapshot, refine: u32) -> Score { + let (w, h) = (clip.width, clip.height); + let mut out = Score::default(); + + assert_eq!( + snap.vectors.len(), + snap.confidence.len(), + "vectors and confidence must carry the same neighbour count and convention" + ); + assert_eq!( + snap.vectors.len(), + clip.truth.len(), + "the snapshot's neighbour count must match the clip's truth, which both index by \ + `neighbour_idx_for_k`" + ); + + for (t, truth) in clip.truth.iter().enumerate() { + let occluded = &clip.occluded[t]; + for ry in 0..refs_along(h) { + for rx in 0..refs_along(w) { + let px = ref_pos(rx, w); + let py = ref_pos(ry, h); + + let mut sum = [0.0f32; 2]; + let mut any_occluded = false; + for y in py..py + PATCH_SIZE { + for x in px..px + PATCH_SIZE { + let idx = (y * w + x) as usize; + sum[0] += truth[idx][0]; + sum[1] += truth[idx][1]; + any_occluded |= occluded[idx]; + } + } + let area = (PATCH_SIZE * PATCH_SIZE) as f32; + let mean = [sum[0] / area, sum[1] / area]; + let mut spread = 0.0f32; + for y in py..py + PATCH_SIZE { + for x in px..px + PATCH_SIZE { + let d = truth[(y * w + x) as usize]; + spread = spread.max((d[0] - mean[0]).abs()).max((d[1] - mean[1]).abs()); + } + } + let kind = if any_occluded { + PatchKind::Occluded + } else if spread > 0.5 { + PatchKind::Boundary + } else { + PatchKind::Plain + }; + + let (bx_lo, bx_hi) = covering_blocks(px, snap.blksize, snap.step, snap.blocks_x); + let (by_lo, by_hi) = covering_blocks(py, snap.blksize, snap.step, snap.blocks_y); + let corner = (by_hi * snap.blocks_x + bx_hi) as usize; + let corner_v = snap.vectors[t][corner]; + let corner_err = endpoint_error(mean, corner_v); + + let mut best_err = corner_err; + for by in by_lo..=by_hi { + for bx in bx_lo..=bx_hi { + let v = snap.vectors[t][(by * snap.blocks_x + bx) as usize]; + best_err = best_err.min(endpoint_error(mean, v)); + } + } + + let k = out.kind_mut(kind); + k.patches += 1; + if corner_err <= refine as f32 { + k.in_window_corner += 1; + } + if best_err <= refine as f32 { + k.in_window_covering += 1; + } + k.epe.push(corner_err); + k.confidence.push(snap.confidence[t][corner]); + } + } + } + + out +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::nl4d::MotionSnapshot; + + /// A 32x32 clip at radius 1 whose truth toward k = +1 is a uniform + /// `[3, 1]`, with nothing occluded. + fn uniform_clip() -> Clip { + let (w, h) = (32u32, 32u32); + let n = (w * h) as usize; + Clip { + width: w, + height: h, + radius: 1, + frames: vec![vec![0.5; n]; 3], + truth: vec![vec![[-3.0, -1.0]; n], vec![[3.0, 1.0]; n]], + occluded: vec![vec![false; n]; 2], + } + } + + /// One vector for every block of every neighbour. + fn uniform_snapshot(vx: i32, vy: i32) -> MotionSnapshot { + let (blocks_x, blocks_y) = (4u32, 4u32); + let blocks = (blocks_x * blocks_y) as usize; + MotionSnapshot { + blocks_x, + blocks_y, + step: 8, + blksize: 16, + offsets: vec![-1, 1], + vectors: vec![vec![[-vx, -vy]; blocks], vec![[vx, vy]; blocks]], + confidence: vec![vec![0.9; blocks]; 2], + } + } + + #[test] + #[should_panic(expected = "neighbour count must match")] + fn score_asserts_the_snapshots_neighbour_count_matches_the_clips_truth() { + // The clip carries truth for both neighbours (k = -1, +1) but + // the snapshot only carries one, so the two disagree on how + // many neighbours `neighbour_idx_for_k` indexes. + let mut snap = uniform_snapshot(3, 1); + snap.vectors.truncate(1); + snap.confidence.truncate(1); + score(&uniform_clip(), &snap, 2); + } + + #[test] + #[should_panic(expected = "same neighbour count and convention")] + fn score_asserts_vectors_and_confidence_carry_the_same_neighbour_count() { + let mut snap = uniform_snapshot(3, 1); + snap.confidence.pop(); + score(&uniform_clip(), &snap, 2); + } + + #[test] + fn covering_blocks_for_the_default_geometry() { + // blksize 16, step 8: patch at 0 is covered by block 0 only, + // patch at 8 by blocks 0 and 1, patch at 16 by blocks 1 and 2. + assert_eq!(covering_blocks(0, 16, 8, 8), (0, 0)); + assert_eq!(covering_blocks(8, 16, 8, 8), (0, 1)); + assert_eq!(covering_blocks(16, 16, 8, 8), (1, 2)); + // step == blksize gives exactly one block. + assert_eq!(covering_blocks(24, 8, 8, 8), (3, 3)); + // The upper end clamps to the grid. + assert_eq!(covering_blocks(56, 16, 8, 7), (6, 6)); + } + + #[test] + fn a_straddling_patch_at_step_equal_blksize_falls_back_to_the_corner_block() { + // blksize 16, step 16: block 0 spans 0..16, block 1 spans + // 16..32. The patch at p = 10 spans 10..18, which no single + // block fully contains. No range is empty here, so the corner + // block (the one the patch's start pixel sits in) is returned + // as the best available search target, matching what the corner + // reading already reads today. + assert_eq!(covering_blocks(10, 16, 16, 8), (0, 0)); + } + + #[test] + fn an_exact_field_scores_every_patch_in_window_with_zero_error() { + let s = score(&uniform_clip(), &uniform_snapshot(3, 1), 2); + assert!(s.plain.patches > 0); + assert_eq!(s.boundary.patches, 0); + assert_eq!(s.occluded.patches, 0); + assert_eq!(s.plain.in_window_rate_corner(), 1.0); + assert_eq!(s.plain.in_window_rate_covering(), 1.0); + assert_eq!(s.plain.epe_mean(), 0.0); + assert!((s.plain.confidence_median() - 0.9).abs() < 1e-6); + } + + #[test] + fn an_error_past_the_refine_window_scores_out_of_window() { + // Off by 3 on x, refine 2: out of window, endpoint error 3. + let s = score(&uniform_clip(), &uniform_snapshot(6, 1), 2); + assert_eq!(s.plain.in_window_rate_corner(), 0.0); + assert!((s.plain.epe_mean() - 3.0).abs() < 1e-6); + assert!((s.plain.epe_p95() - 3.0).abs() < 1e-6); + // Refine 3 admits it. + let s = score(&uniform_clip(), &uniform_snapshot(6, 1), 3); + assert_eq!(s.plain.in_window_rate_corner(), 1.0); + } + + #[test] + fn the_covering_reading_takes_the_best_covering_block() { + // Corner blocks wrong, every other block right. Patches whose + // corner block is wrong but which another block covers still + // count in the covering reading. + let mut snap = uniform_snapshot(3, 1); + for by in 0..4u32 { + for bx in 0..4u32 { + if (bx + by) % 2 == 0 { + snap.vectors[1][(by * 4 + bx) as usize] = [30, 30]; + } + } + } + let s = score(&uniform_clip(), &snap, 2); + assert!(s.plain.in_window_rate_covering() > s.plain.in_window_rate_corner()); + } + + #[test] + fn boundary_and_occluded_patches_are_classified_by_the_truth() { + let mut clip = uniform_clip(); + let w = clip.width as usize; + // A vertical motion boundary at x = 16 toward k = +1. + for y in 0..32usize { + for x in 16..32usize { + clip.truth[1][y * w + x] = [0.0, 0.0]; + } + } + // Pixel column 20 occluded toward k = +1. + for y in 0..32usize { + clip.occluded[1][y * w + 20] = true; + } + let s = score(&clip, &uniform_snapshot(3, 1), 2); + assert!( + s.boundary.patches > 0, + "patches straddling x = 16 are boundary patches" + ); + assert!(s.occluded.patches > 0, "patches touching column 20 are occluded"); + assert!(s.plain.patches > 0); + } +} diff --git a/av-denoise-core/src/nl4d/harness/synth.rs b/av-denoise-core/src/nl4d/harness/synth.rs new file mode 100644 index 0000000..4305d6d --- /dev/null +++ b/av-denoise-core/src/nl4d/harness/synth.rs @@ -0,0 +1,447 @@ +use crate::nlmeans::motion::neighbour_idx_for_k; + +/// A clean luma plane, values in `[0, 1]`. +#[derive(Debug, Clone)] +pub struct Still { + pub width: u32, + pub height: u32, + pub luma: Vec, +} + +impl Still { + /// Parses a binary PGM (`P5`) at 8 or 16 bits per sample. + pub fn from_pgm(bytes: &[u8]) -> Result { + let mut pos = 0usize; + let mut fields: Vec = Vec::new(); + if bytes.len() < 2 || &bytes[..2] != b"P5" { + return Err("not a P5 pgm".to_string()); + } + pos += 2; + while fields.len() < 3 { + while pos < bytes.len() && bytes[pos].is_ascii_whitespace() { + pos += 1; + } + if pos < bytes.len() && bytes[pos] == b'#' { + while pos < bytes.len() && bytes[pos] != b'\n' { + pos += 1; + } + continue; + } + let start = pos; + while pos < bytes.len() && bytes[pos].is_ascii_digit() { + pos += 1; + } + if start == pos { + return Err("malformed pgm header".to_string()); + } + let text = std::str::from_utf8(&bytes[start..pos]).map_err(|e| e.to_string())?; + fields.push(text.parse::().map_err(|e| e.to_string())?); + } + // Exactly one whitespace byte separates maxval from the data. + pos += 1; + let (width, height, maxval) = (fields[0], fields[1], fields[2]); + // Widened to 64 bits so a header claiming huge dimensions cannot + // wrap back into a small `usize` on a 32-bit target and slip + // past the bounds check below. + let n64 = width as u64 * height as u64; + let bytes_per_sample: u64 = if maxval > 255 { 2 } else { 1 }; + let available = (bytes.len() - pos.min(bytes.len())) as u64; + if n64 * bytes_per_sample > available { + return Err(format!( + "pgm data truncated: header claims {width}x{height} at {bytes_per_sample} bytes/sample, \ + only {available} bytes remain" + )); + } + let n = n64 as usize; + let luma = if maxval > 255 { + let data = bytes.get(pos..pos + 2 * n).ok_or("pgm data truncated")?; + data.as_chunks::<2>() + .0 + .iter() + .map(|c| u16::from_be_bytes(*c) as f32 / maxval as f32) + .collect() + } else { + let data = bytes.get(pos..pos + n).ok_or("pgm data truncated")?; + data.iter().map(|&v| v as f32 / maxval as f32).collect() + }; + Ok(Still { width, height, luma }) + } + + /// A textured plane for runs with no real still to hand. + pub fn synthetic(width: u32, height: u32) -> Still { + let mut luma = vec![0.0f32; (width * height) as usize]; + for y in 0..height { + for x in 0..width { + let fx = x as f32 * 0.31; + let fy = y as f32 * 0.23; + let v = 0.5 + 0.2 * (fx.sin() * fy.cos()) + 0.1 * ((fx * 2.7).cos() + (fy * 3.1).sin()); + luma[(y * width + x) as usize] = v.clamp(0.0, 1.0); + } + } + Still { width, height, luma } + } + + /// Samples the still at a fractional position with a Lanczos-3 + /// kernel, clamping to the edge. + fn sample(&self, sx: f32, sy: f32) -> f32 { + const A: i32 = 3; + let lanczos = |t: f32| -> f32 { + if t == 0.0 { + 1.0 + } else if t.abs() >= A as f32 { + 0.0 + } else { + let pt = std::f32::consts::PI * t; + (A as f32 * pt.sin() * (pt / A as f32).sin()) / (pt * pt) + } + }; + let x0 = sx.floor() as i32; + let y0 = sy.floor() as i32; + let mut acc = 0.0f32; + let mut wsum = 0.0f32; + for j in (y0 - A + 1)..=(y0 + A) { + let wy = lanczos(sy - j as f32); + let yy = j.clamp(0, self.height as i32 - 1) as u32; + for i in (x0 - A + 1)..=(x0 + A) { + let w = wy * lanczos(sx - i as f32); + let xx = i.clamp(0, self.width as i32 - 1) as u32; + acc += w * self.luma[(yy * self.width + xx) as usize]; + wsum += w; + } + } + (acc / wsum).clamp(0.0, 1.0) + } +} + +/// The motion each synthetic clip carries. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum MotionClass { + IntegerPan, + HalfPelPan, + Zoom, + CutOut, +} + +/// Scale per frame of the zoom class. +const ZOOM_PER_FRAME: f32 = 1.01; + +impl MotionClass { + pub const ALL: [MotionClass; 4] = [ + MotionClass::IntegerPan, + MotionClass::HalfPelPan, + MotionClass::Zoom, + MotionClass::CutOut, + ]; + + pub fn label(self) -> &'static str { + match self { + MotionClass::IntegerPan => "pan_int", + MotionClass::HalfPelPan => "pan_half", + MotionClass::Zoom => "zoom", + MotionClass::CutOut => "cutout", + } + } + + /// Per-frame velocity of the moving content, in pixels. + pub fn velocity(self) -> [f32; 2] { + match self { + MotionClass::IntegerPan => [3.0, 1.0], + MotionClass::HalfPelPan => [2.5, 0.5], + MotionClass::Zoom => [0.0, 0.0], + MotionClass::CutOut => [4.0, 2.0], + } + } + + /// Top-left corner and side of the cut-out rectangle in the centre + /// frame, a square a third of the shorter side, left of centre so + /// its rightward motion stays inside the frame. + pub fn cut_out_rect(width: u32, height: u32) -> (u32, u32, u32) { + let side = (width.min(height) / 3).max(8); + let x0 = width / 4; + let y0 = (height - side) / 2; + (x0, y0, side) + } +} + +/// A synthetic window of frames with its per-pixel ground truth. +#[derive(Debug, Clone)] +pub struct Clip { + pub width: u32, + pub height: u32, + pub radius: u32, + /// `frames[i]` is the frame at offset `k = i - radius`. + pub frames: Vec>, + /// `truth[t][pixel]` is where the centre frame's pixel lies in + /// neighbour `t`, as a displacement in pixels. + pub truth: Vec>, + /// `occluded[t][pixel]` is true when that pixel has no true match + /// in neighbour `t`. + pub occluded: Vec>, +} + +/// Where the centre frame's pixel `(x, y)` sits in the frame at offset +/// `k`, for the background of `class`. +fn background_displacement(class: MotionClass, k: i32, x: u32, y: u32, width: u32, height: u32) -> [f32; 2] { + match class { + MotionClass::IntegerPan | MotionClass::HalfPelPan => { + let v = class.velocity(); + [v[0] * k as f32, v[1] * k as f32] + }, + MotionClass::Zoom => { + let s = ZOOM_PER_FRAME.powi(k); + let cx = width as f32 / 2.0; + let cy = height as f32 / 2.0; + [(x as f32 - cx) * (s - 1.0), (y as f32 - cy) * (s - 1.0)] + }, + MotionClass::CutOut => [0.0, 0.0], + } +} + +/// Deterministic Gaussian grain from a hashed uniform pair. +fn grain(idx: u32, seed: u32) -> f32 { + let hash = |i: u32| -> f32 { + let mut h = i + .wrapping_mul(2654435761) + .wrapping_add(seed.wrapping_mul(0x9E37_79B9)); + h ^= h >> 15; + h = h.wrapping_mul(0x85EB_CA6B); + h ^= h >> 13; + (h as f32 + 1.0) / (u32::MAX as f32 + 2.0) + }; + let u1 = hash(idx * 2); + let u2 = hash(idx * 2 + 1); + (-2.0 * u1.ln()).sqrt() * (std::f32::consts::TAU * u2).cos() +} + +/// Builds the window of `2 * radius + 1` frames for `class`, with +/// Gaussian grain of `sigma` on every frame, and the ground truth toward +/// every neighbour. +pub fn synthesise(still: &Still, class: MotionClass, radius: u32, sigma: f32, seed: u32) -> Clip { + let (w, h) = (still.width, still.height); + let n = (w * h) as usize; + let (cx0, cy0, side) = MotionClass::cut_out_rect(w, h); + let v = class.velocity(); + + let in_rect_at = |x: f32, y: f32, k: i32| -> bool { + let ox = cx0 as f32 + v[0] * k as f32; + let oy = cy0 as f32 + v[1] * k as f32; + x >= ox && x < ox + side as f32 && y >= oy && y < oy + side as f32 + }; + + let mut frames = Vec::with_capacity((2 * radius + 1) as usize); + for i in 0..(2 * radius + 1) as i32 { + let k = i - radius as i32; + let mut frame = vec![0.0f32; n]; + for y in 0..h { + for x in 0..w { + let idx = (y * w + x) as usize; + let value = if class == MotionClass::CutOut && in_rect_at(x as f32, y as f32, k) { + // The rectangle's content, read from where it sat in the centre. + still.sample(x as f32 - v[0] * k as f32, y as f32 - v[1] * k as f32) + } else { + let d = background_displacement(class, k, x, y, w, h); + // The frame at k shows the centre's pixel p at p + d, so + // pixel (x, y) here comes from the centre's (x, y) - d. + // For a pan and a zoom the inverse map is exact. + match class { + MotionClass::Zoom => { + let s = ZOOM_PER_FRAME.powi(k); + let fx = w as f32 / 2.0 + (x as f32 - w as f32 / 2.0) / s; + let fy = h as f32 / 2.0 + (y as f32 - h as f32 / 2.0) / s; + still.sample(fx, fy) + }, + _ => still.sample(x as f32 - d[0], y as f32 - d[1]), + } + }; + let noise = if sigma > 0.0 { + sigma * grain(idx as u32, seed.wrapping_add(1000 * (i as u32 + 1))) + } else { + 0.0 + }; + frame[idx] = (value + noise).clamp(0.0, 1.0); + } + } + frames.push(frame); + } + + let neighbours = (2 * radius) as usize; + let mut truth = vec![vec![[0.0f32; 2]; n]; neighbours]; + let mut occluded = vec![vec![false; n]; neighbours]; + for k in -(radius as i32)..=(radius as i32) { + if k == 0 { + continue; + } + let t = neighbour_idx_for_k(radius, k) as usize; + for y in 0..h { + for x in 0..w { + let idx = (y * w + x) as usize; + let foreground = class == MotionClass::CutOut && in_rect_at(x as f32, y as f32, 0); + let d = if foreground { + [v[0] * k as f32, v[1] * k as f32] + } else { + background_displacement(class, k, x, y, w, h) + }; + truth[t][idx] = d; + if class == MotionClass::CutOut && !foreground { + occluded[t][idx] = in_rect_at(x as f32 + d[0], y as f32 + d[1], k); + } + } + } + } + + Clip { + width: w, + height: h, + radius, + frames, + truth, + occluded, + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn ramp_still() -> Still { + // Distinct values everywhere, so a shift is visible in any pixel. + let (w, h) = (64u32, 48u32); + let luma = (0..w * h) + .map(|i| ((i % w) as f32 * 0.9 / w as f32 + (i / w) as f32 * 0.1 / h as f32).clamp(0.0, 1.0)) + .collect(); + Still { + width: w, + height: h, + luma, + } + } + + #[test] + fn integer_pan_frames_are_exact_shifts_and_truth_is_the_velocity() { + let still = ramp_still(); + let clip = synthesise(&still, MotionClass::IntegerPan, 1, 0.0, 1); + assert_eq!(clip.frames.len(), 3); + let (w, h) = (clip.width, clip.height); + let v = MotionClass::IntegerPan.velocity(); + // Frame k = +1 holds the still moved by +v. Check an interior pixel. + let (x, y) = (20u32, 20u32); + let moved = clip.frames[2][(y * w + x) as usize]; + let source = + still.luma[((y as i32 - v[1] as i32) as u32 * w + (x as i32 - v[0] as i32) as u32) as usize]; + assert!( + (moved - source).abs() < 1e-6, + "an integer pan must copy pixels exactly" + ); + // Truth toward k = +1 (t = 1 at radius 1) is +v everywhere. + for idx in 0..(w * h) as usize { + assert_eq!(clip.truth[1][idx], v); + assert_eq!(clip.truth[0][idx], [-v[0], -v[1]]); + assert!(!clip.occluded[1][idx]); + } + } + + #[test] + fn half_pel_pan_truth_has_a_half_pixel_component() { + let still = ramp_still(); + let clip = synthesise(&still, MotionClass::HalfPelPan, 1, 0.0, 1); + let v = MotionClass::HalfPelPan.velocity(); + assert!((v[0].fract().abs() - 0.5).abs() < 1e-6 || (v[1].fract().abs() - 0.5).abs() < 1e-6); + assert_eq!(clip.truth[1][100], v); + } + + #[test] + fn zoom_truth_grows_with_distance_from_the_centre() { + let still = ramp_still(); + let clip = synthesise(&still, MotionClass::Zoom, 1, 0.0, 1); + let (w, h) = (clip.width, clip.height); + let centre = ((h / 2) * w + w / 2) as usize; + let corner = 0usize; + let dc = clip.truth[1][centre]; + let dk = clip.truth[1][corner]; + assert!( + dc[0].abs() < 0.01 && dc[1].abs() < 0.01, + "the centre does not move under a zoom" + ); + assert!( + dk[0] < -0.1 && dk[1] < -0.1, + "the top-left corner moves outward, got {dk:?}" + ); + } + + #[test] + fn cut_out_marks_background_hidden_under_the_moved_rectangle() { + let still = ramp_still(); + let clip = synthesise(&still, MotionClass::CutOut, 1, 0.0, 1); + let w = clip.width; + let (x0, y0, side) = MotionClass::cut_out_rect(clip.width, clip.height); + let v = MotionClass::CutOut.velocity(); + // A pixel inside the rectangle in the centre frame moves with it. + let inside = ((y0 + side / 2) * w + x0 + side / 2) as usize; + assert_eq!(clip.truth[1][inside], v); + assert!(!clip.occluded[1][inside]); + // A background pixel just to the right of the rectangle is + // covered once it moves right by v[0] pixels, so toward k = +1 + // it is occluded, and toward k = -1 it is not. + let just_right = ((y0 + side / 2) * w + x0 + side + 1) as usize; + assert_eq!(clip.truth[1][just_right], [0.0, 0.0]); + assert!(clip.occluded[1][just_right]); + assert!(!clip.occluded[0][just_right]); + } + + #[test] + fn grain_has_the_requested_sigma_and_differs_between_frames() { + let still = Still::synthetic(128, 128); + let clip = synthesise(&still, MotionClass::IntegerPan, 1, 0.0, 1); + let noisy = synthesise(&still, MotionClass::IntegerPan, 1, 6.0 / 255.0, 1); + let n = clip.frames[1].len() as f32; + let var: f32 = clip.frames[1] + .iter() + .zip(&noisy.frames[1]) + .map(|(a, b)| (a - b) * (a - b)) + .sum::() + / n; + let sigma = var.sqrt(); + assert!( + (sigma - 6.0 / 255.0).abs() < 0.1 * 6.0 / 255.0, + "measured sigma {sigma}" + ); + assert_ne!( + noisy.frames[0], noisy.frames[1], + "each frame carries its own grain" + ); + } + + #[test] + fn pgm_parses_8_and_16_bit_planes() { + let mut p8 = b"P5\n# comment\n2 2\n255\n".to_vec(); + p8.extend_from_slice(&[0, 128, 255, 64]); + let s = Still::from_pgm(&p8).expect("8-bit parse"); + assert_eq!((s.width, s.height), (2, 2)); + assert!((s.luma[1] - 128.0 / 255.0).abs() < 1e-6); + let mut p16 = b"P5 2 1 65535\n".to_vec(); + p16.extend_from_slice(&[0xFF, 0xFF, 0x00, 0x00]); + let s = Still::from_pgm(&p16).expect("16-bit parse"); + assert_eq!(s.luma, vec![1.0, 0.0]); + } + + #[test] + fn a_header_claiming_more_data_than_is_present_is_rejected() { + // The header claims a 100x100 plane, but only 4 sample bytes + // follow. Also exercises the overflow-safe path: `width * + // height` for a header this large already overflows `u32`. + let mut p8 = b"P5\n100 100\n255\n".to_vec(); + p8.extend_from_slice(&[0, 128, 255, 64]); + let err = Still::from_pgm(&p8).expect_err("truncated data must be rejected"); + assert!(err.contains("truncated"), "got {err}"); + } + + #[test] + fn a_header_with_dimensions_that_overflow_u32_is_rejected_not_wrapped() { + // `width * height` overflows `u32` here; widening to `u64` + // before multiplying must catch this as truncated data rather + // than wrapping to a small value that a short buffer satisfies. + let p8 = b"P5\n70000 70000\n255\n".to_vec(); + let err = Still::from_pgm(&p8).expect_err("an overflowing header must be rejected"); + assert!(err.contains("truncated"), "got {err}"); + } +} diff --git a/av-denoise-core/src/nl4d/kernels/mod.rs b/av-denoise-core/src/nl4d/kernels/mod.rs new file mode 100644 index 0000000..94c1670 --- /dev/null +++ b/av-denoise-core/src/nl4d/kernels/mod.rs @@ -0,0 +1,7 @@ +//! GPU kernels that belong to nl4d alone. + +#![doc(hidden)] + +mod regularise; + +pub use regularise::nl4d_mv_regularise; diff --git a/av-denoise-core/src/nl4d/kernels/regularise.rs b/av-denoise-core/src/nl4d/kernels/regularise.rs new file mode 100644 index 0000000..1cc1abf --- /dev/null +++ b/av-denoise-core/src/nl4d/kernels/regularise.rs @@ -0,0 +1,237 @@ +use cubecl::prelude::*; + +/// How many vectors a block considers, its own, the neighbourhood +/// median, the four adjacent blocks' and zero. +pub const REGULARISE_CANDIDATES: u32 = 7; + +/// The most neighbours a block has in its 3x3 neighbourhood. +const NEIGHBOURHOOD: u32 = 8; + +#[cube] +fn clamp_coord(value: i32, limit: i32) -> i32 { + let mut result = value; + if value < 0 { + result = 0; + } else if value >= limit { + result = limit - 1; + } + result +} + +#[cube] +fn abs_i32(value: i32) -> i32 { + let mut result = value; + if value < 0 { + result = -value; + } + result +} + +/// Sorts the first `n` entries of `vals` in place and returns the lower +/// median. +#[cube] +fn median_of(vals: &mut Array, n: u32) -> i32 { + let mut i: u32 = 1; + while i < n { + let key = vals[i as usize]; + let mut j = i; + while j > 0u32 && vals[(j - 1u32) as usize] > key { + vals[j as usize] = vals[(j - 1u32) as usize]; + j -= 1u32; + } + vals[j as usize] = key; + i += 1u32; + } + vals[((n - 1u32) / 2u32) as usize] +} + +/// Re-scores one block's motion vector against its neighbourhood and +/// writes the winner, with a fresh confidence, to the output field. +/// +/// One cube handles one block of the field. Thread 0 gathers the 3x3 +/// neighbourhood's vectors from `mv_in`, takes their component-wise +/// median, and lays out the candidates in shared memory. The block's +/// own vector is candidate 0. Each of the next threads scores one +/// candidate by SAD over the block on the level-0 luma planes, plus +/// `lambda_pixel` times the candidate's distance from the median in +/// pixels. Thread 0 then picks the lowest cost, and a tie keeps the +/// earlier candidate, so the block's own vector wins every tie. +/// +/// The winner's confidence is derived from its SAD exactly as +/// `nlm_mc_block_match_fine` derives it, with the same +/// `sad_noise_floor` and `thsad`. +/// +/// `mv_in` and `mv_out` are separate buffers. Every block reads the +/// whole input before any block's output exists, so the result does +/// not depend on block order. +#[cube(launch_unchecked)] +#[expect( + clippy::too_many_arguments, + reason = "every argument is a buffer or comptime shape the kernel binds" +)] +pub fn nl4d_mv_regularise( + centre: &Array, + neighbour: &Array, + mv_in: &Array, + mv_out: &mut Array, + confidence_out: &mut Array, + lambda_pixel: f32, + sad_noise_floor: f32, + thsad: f32, + #[comptime] width: u32, + #[comptime] height: u32, + #[comptime] blksize: u32, + #[comptime] step: u32, + #[comptime] blocks_x: u32, + #[comptime] blocks_y: u32, +) { + let bx = CUBE_POS_X; + let by = CUBE_POS_Y; + let block = by * blocks_x + bx; + let local_x = UNIT_POS_X; + let local_y = UNIT_POS_Y; + let thread_id = local_y * CUBE_DIM_X + local_x; + + let block_pixels = comptime!(blksize * blksize); + let mut centre_smem = SharedMemory::::new(block_pixels as usize); + let mut cand = SharedMemory::::new(comptime!(2 * REGULARISE_CANDIDATES) as usize); + let mut median = SharedMemory::::new(2usize); + let mut sad_scratch = SharedMemory::::new(REGULARISE_CANDIDATES as usize); + let mut cost = SharedMemory::::new(REGULARISE_CANDIDATES as usize); + + let block_origin_x = bx as i32 * step as i32; + let block_origin_y = by as i32 * step as i32; + + // The centre tile, loaded once and shared by every candidate. + let mut py = local_y; + while py < blksize { + let mut px = local_x; + while px < blksize { + let cx = clamp_coord(block_origin_x + px as i32, width as i32); + let cy = clamp_coord(block_origin_y + py as i32, height as i32); + centre_smem[(py * blksize + px) as usize] = centre[(cy * width as i32 + cx) as usize]; + px += CUBE_DIM_X; + } + py += CUBE_DIM_Y; + } + + if thread_id == 0u32 { + let mut xs = Array::::new(NEIGHBOURHOOD as usize); + let mut ys = Array::::new(NEIGHBOURHOOD as usize); + let mut n: u32 = 0; + let mut dy: u32 = 0; + while dy < 3u32 { + let mut dx: u32 = 0; + while dx < 3u32 { + if dx != 1u32 || dy != 1u32 { + let nx = bx as i32 + dx as i32 - 1i32; + let ny = by as i32 + dy as i32 - 1i32; + if nx >= 0 && ny >= 0 && nx < blocks_x as i32 && ny < blocks_y as i32 { + let idx = ((ny as u32 * blocks_x + nx as u32) * 2u32) as usize; + xs[n as usize] = mv_in[idx]; + ys[n as usize] = mv_in[idx + 1]; + n += 1u32; + } + } + dx += 1u32; + } + dy += 1u32; + } + let own_x = mv_in[(block * 2u32) as usize]; + let own_y = mv_in[(block * 2u32 + 1u32) as usize]; + // A block with no neighbours is its own median. + let mut mx = own_x; + let mut my = own_y; + if n > 0u32 { + mx = median_of(&mut xs, n); + my = median_of(&mut ys, n); + } + median[0] = mx; + median[1] = my; + + cand[0] = own_x; + cand[1] = own_y; + cand[2] = mx; + cand[3] = my; + // Left, right, up, down. Off-grid neighbours repeat the block's + // own vector, which the tie rule then discards. + let mut c: u32 = 2; + let mut side: u32 = 0; + while side < 4u32 { + let mut nx = bx as i32; + let mut ny = by as i32; + if side == 0u32 { + nx -= 1; + } else if side == 1u32 { + nx += 1; + } else if side == 2u32 { + ny -= 1; + } else { + ny += 1; + } + let mut vx = own_x; + let mut vy = own_y; + if nx >= 0 && ny >= 0 && nx < blocks_x as i32 && ny < blocks_y as i32 { + let idx = ((ny as u32 * blocks_x + nx as u32) * 2u32) as usize; + vx = mv_in[idx]; + vy = mv_in[idx + 1]; + } + cand[(c * 2u32) as usize] = vx; + cand[(c * 2u32 + 1u32) as usize] = vy; + c += 1u32; + side += 1u32; + } + cand[(c * 2u32) as usize] = 0; + cand[(c * 2u32 + 1u32) as usize] = 0; + } + sync_cube(); + + if thread_id < REGULARISE_CANDIDATES { + let mvx = cand[(thread_id * 2u32) as usize]; + let mvy = cand[(thread_id * 2u32 + 1u32) as usize]; + let mut sad: f32 = 0.0; + for iy in 0..blksize { + for ix in 0..blksize { + let cx = block_origin_x + ix as i32; + let cy = block_origin_y + iy as i32; + let centre_val = centre_smem[(iy * blksize + ix) as usize]; + let nx = clamp_coord(cx + mvx, width as i32); + let ny = clamp_coord(cy + mvy, height as i32); + let diff = centre_val - neighbour[(ny * width as i32 + nx) as usize]; + let abs_diff = if diff < 0.0f32 { -diff } else { diff }; + sad += abs_diff; + } + } + let deviation = abs_i32(mvx - median[0]) + abs_i32(mvy - median[1]); + sad_scratch[thread_id as usize] = sad; + cost[thread_id as usize] = sad + lambda_pixel * deviation as f32; + } + sync_cube(); + + if thread_id == 0u32 { + let mut best: u32 = 0; + let mut best_cost = cost[0]; + let mut c: u32 = 1; + while c < REGULARISE_CANDIDATES { + if cost[c as usize] < best_cost { + best_cost = cost[c as usize]; + best = c; + } + c += 1u32; + } + mv_out[(block * 2u32) as usize] = cand[(best * 2u32) as usize]; + mv_out[(block * 2u32 + 1u32) as usize] = cand[(best * 2u32 + 1u32) as usize]; + + let mut excess = sad_scratch[best as usize] - sad_noise_floor; + if excess < 0.0f32 { + excess = 0.0f32; + } + let thsad_sq = thsad * thsad; + let excess_sq = excess * excess; + let mut confidence = (thsad_sq - excess_sq) / (thsad_sq + excess_sq); + if confidence < 0.0f32 { + confidence = 0.0f32; + } + confidence_out[block as usize] = confidence; + } +} diff --git a/av-denoise-core/src/nl4d/mod.rs b/av-denoise-core/src/nl4d/mod.rs index 3d80e4b..c8dcdaa 100644 --- a/av-denoise-core/src/nl4d/mod.rs +++ b/av-denoise-core/src/nl4d/mod.rs @@ -13,7 +13,11 @@ //! a single-frame search can. mod denoiser; +pub mod harness; +pub mod kernels; mod params; +mod regularise; +mod snapshot; // Every test in this tree runs against a real GPU runtime, see // `tests::helpers::R`, so it only builds when a wgpu-backed feature is @@ -23,3 +27,4 @@ mod tests; pub use denoiser::Nl4dDenoiser; pub use params::{MAX_KAISER_BETA, MAX_MISMATCH_SCALE, Nl4dParams}; +pub use snapshot::MotionSnapshot; diff --git a/av-denoise-core/src/nl4d/params.rs b/av-denoise-core/src/nl4d/params.rs index 87344a0..b6d5d39 100644 --- a/av-denoise-core/src/nl4d/params.rs +++ b/av-denoise-core/src/nl4d/params.rs @@ -2,12 +2,14 @@ use crate::nlmeans::{ChannelMode, HqParams, MotionCompensationMode, MotionEstima /// The largest [`Nl4dParams::mismatch_scale`] worth accepting. /// -/// The mismatch variance is capped at -/// [`crate::collab::kernels::fused::MEMBER_SIGMA2_CAP`] times the channel -/// variance, and the worst-matched blocks reach that cap at a scale of -/// roughly `319 * sigma`. Even a source noisy enough to measure `sigma = -/// 0.05` saturates below 16, so nothing above this can move a pixel and -/// accepting it would only promise a range that is not there. +/// A member's own match distance never exceeds `3 * PATCH_AREA` in the +/// search's units, so its mismatch variance never exceeds +/// `mismatch_scale^2` in absolute pixel-value units. The mechanism caps +/// at [`crate::collab::kernels::fused::MEMBER_SIGMA2_CAP`] times the +/// channel variance, so even the worst possible mismatch saturates by a +/// scale of `8 * sigma`. Even a source noisy enough to measure `sigma = +/// 0.05` saturates well under 1, so nothing above this can move a pixel +/// and accepting it would only promise a range that is not there. pub const MAX_MISMATCH_SCALE: f32 = 16.0; /// The largest [`Nl4dParams::kaiser_beta`] worth accepting. @@ -20,6 +22,20 @@ pub const MAX_MISMATCH_SCALE: f32 = 16.0; /// weights fall under what the fixed-point accumulators resolve. pub const MAX_KAISER_BETA: f32 = 8.0; +/// The most motion blocks that may cover a reference patch on one axis. +/// +/// A block grid at a step below `blksize` puts several blocks over one +/// patch, and +/// [`crate::collab::kernels::fused::collab_fused`] searches all of them. +/// It unrolls its per-neighbour duplicate-rectangle arrays over the +/// square of this bound, so the bound is what caps the shader's register +/// footprint. At 4 the arrays hold 16 rectangles and the shipped +/// geometry, `blksize = 16` at `overlap = 8`, uses 2. +/// +/// The step is `blksize - overlap`, so 4 admits an overlap of up to +/// three quarters of the block size. +pub const MAX_COVERING_BLOCKS: u32 = 4; + /// Tuning for [`super::Nl4dDenoiser`]. /// /// `nlm` supplies the front end that builds the frame ring, the motion @@ -48,27 +64,28 @@ pub struct Nl4dParams { /// Higher shrinks more coefficients, so it removes more noise and /// more fine detail. /// - /// Defaults to 5.3, the luma value. Chroma wants a different one, and - /// callers building `Nl4dParams` directly get no per-plane - /// resolution. See [`crate::nl4d_default_lambda_ht`]. + /// Defaults to 5.2. Note that in reality luma and chroma want separately + /// tuned values. See [nl4d_default_lambda_ht](crate::nl4d_default_lambda_ht). pub lambda_ht: f32, /// The confidence floor below which a whole neighbour block is /// skipped rather than scored, in `[0, 1)`. Only affects how much /// compute a submit spends, never which candidates are admitted once /// they are scored. pub c_min: f32, - /// A multiplier on the mismatch variance a poorly matched temporal - /// member carries into the hard threshold. + /// A multiplier on the mismatch variance a temporal member carries + /// into the hard threshold. /// - /// The variance grows with the square of this, so `2.0` is a - /// four-fold increase. `1.0`, the default, is the shipped + /// A member matched in a neighbour frame is treated as a noisier + /// observation of the reference, and its extra variance is its own + /// match distance, per channel and per pixel, with the noise floor + /// removed. The variance grows with the square of this, so `2.0` is + /// a four-fold increase. `1.0`, the default, is the shipped /// calibration. `0.0` matches `confidence_variance: false`. /// /// The mechanism saturates. A member's extra variance is capped at /// [`crate::collab::kernels::fused::MEMBER_SIGMA2_CAP`] times the - /// channel variance, which the worst-matched blocks reach somewhere - /// between 3 and 13 depending on how noisy the source is, so raising - /// this past that point stops changing anything. + /// channel variance, so raising this past the point where a + /// member's distance reaches the cap stops changing anything. pub mismatch_scale: f32, /// The `beta` of the Kaiser window each filtered patch is tapered /// with as it is aggregated, in `0..=8`. @@ -90,6 +107,19 @@ pub struct Nl4dParams { /// member the plain channel sigma instead, which is what an ablation /// needs to isolate the effect of this mechanism. pub confidence_variance: bool, + /// The penalty on a block's vector deviating from its + /// neighbourhood's median, in the field regularisation pass. + /// + /// The pass re-scores each block's vector against the median of its + /// neighbours, the four adjacent blocks' vectors and zero, adding + /// this times the distance from the median, in pixels, scaled so + /// `1.0` weighs one pixel of deviation like a 5/255 per-pixel + /// mismatch. Defaults to `1.0`, calibrated with a `field_lambda` + /// sweep on the `mc_accuracy` bench. The pass gains most of its + /// accuracy by a moderate penalty and further increases add little, + /// so `1.0` sits inside that plateau rather than at its edge. `0.0` + /// skips the pass. + pub field_lambda: f32, } impl Default for Nl4dParams { @@ -111,11 +141,12 @@ impl Default for Nl4dParams { temporal_radius: 2, refine: 2, spatial_radius: 9, - lambda_ht: 5.3, + lambda_ht: 5.2, c_min: 0.05, mismatch_scale: 1.0, kaiser_beta: 2.0, confidence_variance: true, + field_lambda: 1.0, } } } @@ -149,6 +180,26 @@ impl Nl4dParams { ); } + // Only checked once the geometry itself is sound. An overlap at + // or past blksize gives a step of 0, which `nlm.validate()` + // rejects on its own terms below with the real fault named. Left + // unguarded, that same case saturates the step to 1 here and + // reports a nonsensical covering-block count instead. + if let MotionCompensationMode::Mvtools { blksize, overlap, .. } = self.nlm.motion_compensation + && overlap < blksize + { + let step = blksize - overlap; + let covers = blksize.div_ceil(step); + if covers > MAX_COVERING_BLOCKS { + return Err(format!( + "nlm.motion_compensation blksize={blksize} at overlap={overlap} gives a step \ + of {step}, so {covers} blocks cover a patch on each axis, past the \ + {MAX_COVERING_BLOCKS} the temporal grouping kernel unrolls its search over. \ + Raise the step by lowering the overlap." + )); + } + } + if !(1..=crate::collab::MAX_TEMPORAL_RADIUS).contains(&self.temporal_radius) { return Err(format!( "temporal_radius={} must be in 1..={}", @@ -193,6 +244,13 @@ impl Nl4dParams { )); } + if !(self.field_lambda.is_finite() && self.field_lambda >= 0.0) { + return Err(format!( + "field_lambda must be finite and at least 0, got {}", + self.field_lambda + )); + } + Ok(()) } } @@ -240,6 +298,106 @@ mod tests { } } + /// A block geometry with `blksize / step` at or under + /// [`MAX_COVERING_BLOCKS`] is what the grouping kernel unrolls its + /// search over. + /// + /// The shipped geometry gives a step of 8 and so 2 covering blocks. + /// An overlap of three quarters of the block size gives a step of 4 + /// and exactly 4, the boundary. + #[test] + fn validate_accepts_block_geometries_up_to_the_covering_bound() { + for (blksize, overlap, covers) in [(16u32, 8u32, 2u32), (16, 12, 4), (32, 24, 4), (8, 4, 2)] { + let params = Nl4dParams { + nlm: NlmParams { + motion_compensation: MotionCompensationMode::Mvtools { + blksize, + overlap, + search_radius: 4, + pyramid_levels: 2, + estimation: MotionEstimation::Auto, + }, + ..Nl4dParams::default().nlm + }, + ..Nl4dParams::default() + }; + assert!( + params.validate().is_ok(), + "blksize={blksize} overlap={overlap} covers {covers} blocks and should be accepted" + ); + } + } + + /// Past the bound the kernel would unroll a far larger duplicate + /// check and hold far more rectangles in registers, so the + /// configuration is refused rather than compiled. + #[test] + fn validate_rejects_a_block_geometry_past_the_covering_bound() { + for (blksize, overlap) in [(16u32, 13u32), (16, 14), (32, 31), (32, 25)] { + let params = Nl4dParams { + nlm: NlmParams { + motion_compensation: MotionCompensationMode::Mvtools { + blksize, + overlap, + search_radius: 4, + pyramid_levels: 2, + estimation: MotionEstimation::Auto, + }, + ..Nl4dParams::default().nlm + }, + ..Nl4dParams::default() + }; + let err = params + .validate() + .expect_err("a step this small should be rejected"); + assert!( + err.contains(&format!("blksize={blksize}")) && err.contains(&format!("overlap={overlap}")), + "error should name the offending blksize and overlap, got {err}" + ); + } + } + + /// An overlap equal to blksize gives a step of 0, which is really a + /// `nlm.validate()` fault, not a covering-block one. `Nl4dParams`'s + /// own check has to stay quiet about it, mirroring how construction + /// runs both validations in sequence, so the caller sees the overlap + /// constraint named rather than a nonsensical covering-block count + /// computed from a saturated step. + #[test] + fn overlap_equal_to_blksize_reports_the_overlap_constraint_not_covering_blocks() { + let params = Nl4dParams { + nlm: NlmParams { + motion_compensation: MotionCompensationMode::Mvtools { + blksize: 16, + overlap: 16, + search_radius: 4, + pyramid_levels: 2, + estimation: MotionEstimation::Auto, + }, + ..Nl4dParams::default().nlm + }, + ..Nl4dParams::default() + }; + assert!( + params.validate().is_ok(), + "the covering-block check must not fire on a geometry nlm.validate() rejects on its \ + own terms" + ); + let err = params + .nlm + .validate() + .expect_err("overlap == blksize must be rejected") + .to_string(); + assert!( + err.contains("overlap") && err.contains("blksize"), + "error should name the overlap constraint, got {err}" + ); + assert!( + !err.contains("cover a patch"), + "error should not be the covering-block message, got {err}" + ); + } + #[test] fn validate_rejects_missing_hq() { let params = Nl4dParams { @@ -353,4 +511,35 @@ mod tests { assert!(params.validate().is_err(), "c_min={bad} should be rejected"); } } + + #[test] + fn validate_accepts_zero_and_positive_field_lambda() { + for lambda in [0.0, 0.5, 4.0] { + let params = Nl4dParams { + field_lambda: lambda, + ..Nl4dParams::default() + }; + assert!( + params.validate().is_ok(), + "field_lambda={lambda} should be accepted" + ); + } + } + + #[test] + fn validate_rejects_negative_or_non_finite_field_lambda() { + for lambda in [-0.1, f32::NAN, f32::INFINITY] { + let params = Nl4dParams { + field_lambda: lambda, + ..Nl4dParams::default() + }; + let err = params + .validate() + .expect_err("field_lambda={lambda} should be rejected"); + assert!( + err.contains("field_lambda"), + "error should name field_lambda, got {err}" + ); + } + } } diff --git a/av-denoise-core/src/nl4d/regularise.rs b/av-denoise-core/src/nl4d/regularise.rs new file mode 100644 index 0000000..cb5edf2 --- /dev/null +++ b/av-denoise-core/src/nl4d/regularise.rs @@ -0,0 +1,85 @@ +use cubecl::prelude::*; +use cubecl::server::Handle; + +use super::kernels::nl4d_mv_regularise; +use crate::nlmeans::RingView; +use crate::nlmeans::motion::{ + MotionCtx, + THSAD_PIXEL, + confidence_byte_offset, + level_dims, + mv_field_byte_offset, + pyramid_slot_byte_offset, +}; + +/// Runs the field regularisation pass over every neighbour of `view`, +/// writing the result into `mv_out` and `conf_out`, which share the +/// front end's per-neighbour layout. +#[expect( + clippy::too_many_arguments, + reason = "the dispatch threads through every buffer and shape the kernel binds" +)] +pub(super) fn run_regularise( + client: &ComputeClient, + mc: &MotionCtx, + view: &RingView, + width: u32, + height: u32, + field_lambda: f32, + sad_noise_floor: f32, + thsad: f32, + mv_out: &Handle, + conf_out: &Handle, +) -> Result<(), anyhow::Error> { + let (fw, fh) = level_dims(width, height, 0); + let level_len = (fw * fh) as usize; + let blocks = (mc.blocks_x * mc.blocks_y) as usize; + let lambda_pixel = field_lambda * (mc.blksize * mc.blksize) as f32 * THSAD_PIXEL; + let centre = view.pyramid.clone().offset_start(pyramid_slot_byte_offset( + width, + height, + view.frame_count, + 0, + view.centre_slot, + mc.align, + )); + + for (t, &slot) in view.neighbour_slots.iter().enumerate() { + let t = t as u32; + let neighbour = view.pyramid.clone().offset_start(pyramid_slot_byte_offset( + width, + height, + view.frame_count, + 0, + slot, + mc.align, + )); + let mv_in = view.mv_field.clone().offset_start(mv_field_byte_offset(mc, t)); + let mv_dst = mv_out.clone().offset_start(mv_field_byte_offset(mc, t)); + let conf_dst = conf_out.clone().offset_start(confidence_byte_offset(mc, t)); + + unsafe { + nl4d_mv_regularise::launch_unchecked::( + client, + CubeCount::new_2d(mc.blocks_x, mc.blocks_y), + CubeDim::new_2d(8, 8), + ArrayArg::from_raw_parts(centre.clone(), level_len), + ArrayArg::from_raw_parts(neighbour, level_len), + ArrayArg::from_raw_parts(mv_in, 2 * blocks), + ArrayArg::from_raw_parts(mv_dst, 2 * blocks), + ArrayArg::from_raw_parts(conf_dst, blocks), + lambda_pixel, + sad_noise_floor, + thsad, + fw, + fh, + mc.blksize, + mc.step, + mc.blocks_x, + mc.blocks_y, + ); + } + } + + Ok(()) +} diff --git a/av-denoise-core/src/nl4d/snapshot.rs b/av-denoise-core/src/nl4d/snapshot.rs new file mode 100644 index 0000000..68ce161 --- /dev/null +++ b/av-denoise-core/src/nl4d/snapshot.rs @@ -0,0 +1,88 @@ +use cubecl::prelude::*; +use cubecl::server::Handle; + +/// The motion field and confidence the last collaborative pass read, +/// copied back to the host. +/// +/// `vectors[t][block]` is block `block`'s vector toward neighbour `t`, +/// in pixels, and `confidence[t][block]` that block's confidence in +/// `[0, 1]`. `offsets[t]` is neighbour `t`'s temporal offset from the +/// centre frame. Blocks run row-major over `blocks_x * blocks_y`, and +/// block `(bx, by)` covers `blksize` pixels starting at `bx * step` on +/// each axis. +/// +/// This exists for measurement tooling. It is not a stable interface. +#[doc(hidden)] +#[derive(Debug, Clone, PartialEq)] +pub struct MotionSnapshot { + pub blocks_x: u32, + pub blocks_y: u32, + pub step: u32, + pub blksize: u32, + pub offsets: Vec, + pub vectors: Vec>, + pub confidence: Vec>, +} + +/// The device buffers one pass handed the fused kernel, kept so the +/// snapshot can read them back after the fact. +pub(super) struct LastFields { + pub mv_field: Handle, + pub confidence: Handle, + pub mv_stride: u32, + pub conf_stride: u32, + pub neighbours: u32, +} + +/// Reads `fields` back and unpacks each neighbour's slice. +pub(super) fn read_snapshot( + client: &ComputeClient, + fields: &LastFields, + radius: u32, + blocks_x: u32, + blocks_y: u32, + step: u32, + blksize: u32, +) -> MotionSnapshot { + let blocks = (blocks_x * blocks_y) as usize; + let mv_bytes = client + .read_one(fields.mv_field.clone()) + .expect("motion field readback failed"); + let mv = i32::from_bytes(&mv_bytes); + let conf_bytes = client + .read_one(fields.confidence.clone()) + .expect("confidence readback failed"); + let conf = f32::from_bytes(&conf_bytes); + + let mut offsets = Vec::with_capacity(fields.neighbours as usize); + let mut vectors = Vec::with_capacity(fields.neighbours as usize); + let mut confidence = Vec::with_capacity(fields.neighbours as usize); + for t in 0..fields.neighbours { + // Mirrors `neighbour_idx_for_k`, ascending k on the negative + // side first, then ascending k on the positive side. + let k = if t < radius { + t as i32 - radius as i32 + } else { + t as i32 - radius as i32 + 1 + }; + offsets.push(k); + let mv_base = (t * fields.mv_stride) as usize; + vectors.push( + (0..blocks) + .map(|b| [mv[mv_base + 2 * b], mv[mv_base + 2 * b + 1]]) + .collect(), + ); + let c_base = (t * fields.conf_stride) as usize; + confidence.push(conf[c_base..c_base + blocks].to_vec()); + } + + MotionSnapshot { + blocks_x, + blocks_y, + step, + blksize, + offsets, + vectors, + confidence, + } +} diff --git a/av-denoise-core/src/nl4d/tests/confidence.rs b/av-denoise-core/src/nl4d/tests/confidence.rs index 82ed06f..fc3cec2 100644 --- a/av-denoise-core/src/nl4d/tests/confidence.rs +++ b/av-denoise-core/src/nl4d/tests/confidence.rs @@ -1,120 +1,12 @@ use cubecl::prelude::*; -use super::grouping::{BLKSIZE, THSAD}; use super::helpers::{R, make_client, noisy_copy_of, textured_base}; -use crate::collab::kernels::fused::mismatch_sigma2; use crate::collab::kernels::transforms::haar_variance_ladder; use crate::nl4d::{Nl4dDenoiser, Nl4dParams}; use crate::nlmeans::{ChannelMode, HqParams, MotionCompensationMode, MotionEstimation, NlmParams}; const REFINE: u32 = 2; -/// The formula [`mismatch_sigma2`] runs on the GPU, mirrored on the host -/// for these tests, with the same argument order and the same -/// operations, so floating-point rounding matches to well within the -/// tolerances below. -fn expected_mismatch_sigma2(confidence: f32, thsad: f32, blksize: u32) -> f32 { - let blksize_area = (blksize * blksize) as f32; - let ratio = (1.0 - confidence) / (1.0 + confidence); - let e2 = thsad * thsad * ratio; - let eps = e2.sqrt() / blksize_area; - std::f32::consts::FRAC_PI_2 * eps * eps -} - -/// Runs [`mismatch_sigma2`] on the GPU, one confidence per thread, so -/// the host mirror above is checked against the code the filter -/// actually calls rather than against itself. -#[cube(launch_unchecked)] -fn mismatch_sigma2_probe( - confidence: &Array, - thsad: f32, - blksize_area: f32, - out: &mut Array, - #[comptime] n: u32, -) { - let i = ABSOLUTE_POS_X; - if i < n { - out[i as usize] = mismatch_sigma2(confidence[i as usize], thsad, blksize_area); - } -} - -fn run_mismatch_sigma2(confidences: &[f32], thsad: f32, blksize: u32) -> Vec { - let client = make_client(); - let n = confidences.len(); - let conf_buf = client.create_from_slice(f32::as_bytes(confidences)); - // One output slot per confidence. `size_of_val(confidences)` reaches - // the same number but ties the output's size to the input's slice. - #[expect( - clippy::manual_slice_size_calculation, - reason = "n is the element count this output holds, not the input's byte length" - )] - let out_buf = client.empty(n * size_of::()); - - unsafe { - mismatch_sigma2_probe::launch_unchecked::( - &client, - CubeCount::new_1d(1), - CubeDim::new_1d(64), - ArrayArg::from_raw_parts(conf_buf, n), - thsad, - (blksize * blksize) as f32, - ArrayArg::from_raw_parts(out_buf.clone(), n), - n as u32, - ); - } - - let bytes = client.read_one(out_buf).expect("mismatch_sigma2 readback failed"); - f32::from_bytes(&bytes)[..n].to_vec() -} - -/// `c = 1.0`, a perfect motion match, must give `sigma_m2 = 0.0` -/// exactly, whatever `thsad` is. -#[test] -fn confidence_one_gives_zero_mismatch_variance() { - for thsad in [THSAD, 0.5, 12.0] { - let out = run_mismatch_sigma2(&[1.0f32], thsad, BLKSIZE); - assert_eq!( - out[0], 0.0, - "a perfect match must carry no mismatch variance at thsad={thsad}, got {}", - out[0] - ); - } -} - -/// A known confidence must produce exactly the variance -/// [`expected_mismatch_sigma2`] derives, over the whole range the -/// confidence field can hold. -#[test] -fn low_confidence_produces_the_derived_mismatch_variance() { - let confidences = [0.0f32, 0.05, 0.2, 0.5, 0.8, 0.95]; - let out = run_mismatch_sigma2(&confidences, THSAD, BLKSIZE); - - for (idx, &c) in confidences.iter().enumerate() { - let expected = expected_mismatch_sigma2(c, THSAD, BLKSIZE); - assert!( - (out[idx] - expected).abs() < 1e-9, - "expected sigma_m2 {expected} for confidence {c}, got {}", - out[idx] - ); - } - - // Sanity: a low confidence must actually derive a value far from - // zero, or the assertions above would pass trivially against a - // broken formula that always returns ~0. - let low = expected_mismatch_sigma2(0.2, THSAD, BLKSIZE); - assert!(low > 1e-4, "expected a non-trivial mismatch variance, got {low}"); -} - -// The mismatch variance only ever reaches a motion-predicted member. -// A centre-frame member is not motion-predicted, so there is no -// mismatch to model and it keeps the plain `sigma^2` whatever the -// confidence field holds. -// `collab::tests::fused::centre_frame_members_ignore_the_confidence_field` -// runs the filter with every neighbour gated out, leaving nothing but -// centre-frame members, and shows the `confidence_variance` flag then -// changes nothing at all even with the confidence buffer at 0.0, the -// worst value the formula above can see. - /// Inflating exactly one member's variance must raise exactly the stack /// rows that member participates in and leave every other row /// unchanged, checked against the host-mirror ladder @@ -202,6 +94,7 @@ fn mismatch_scale_test_params( // The shipped default, so these run the aggregation a real // caller gets. kaiser_beta: 2.0, + field_lambda: 0.0, } } diff --git a/av-denoise-core/src/nl4d/tests/grouping.rs b/av-denoise-core/src/nl4d/tests/grouping.rs index 6dfca20..ee4bf64 100644 --- a/av-denoise-core/src/nl4d/tests/grouping.rs +++ b/av-denoise-core/src/nl4d/tests/grouping.rs @@ -13,19 +13,14 @@ use crate::collab::geometry::{fused_cubes_x, ref_count, refs_along}; use crate::collab::kernels::aggregate::{cross_frame_accum_scale, kaiser_window, weight_scale}; use crate::collab::kernels::fused::collab_fused; use crate::collab::kernels::transforms::dct_noise_profile; -use crate::collab::{PATCH_SIZE, STEP}; +use crate::collab::{PATCH_AREA, PATCH_SIZE, STEP}; -/// The motion block side length these fixtures score confidence and -/// mismatch variance against, distinct from [`BLK_STEP`], which stays -/// at `PATCH_SIZE` so a block boundary lines up with a patch boundary. +/// The motion block side length these fixtures score confidence +/// against, distinct from [`BLK_STEP`], which stays at `PATCH_SIZE` so +/// a block boundary lines up with a patch boundary. The mismatch +/// variance is scored against a member's own match distance instead. pub(super) const BLKSIZE: u32 = 16; -/// `thsad(BLKSIZE, 1.0)` in normalised SAD units, the same value real -/// callers get from `NlmDenoiser::thsad_value` at this block size and -/// the default `thsad_scale`. Duplicated here rather than imported, -/// since `motion::thsad` is crate-private to `nlmeans`. -pub(super) const THSAD: f32 = (BLKSIZE * BLKSIZE) as f32 * 0.02; - const REFINE: u32 = 2; const K_MAX: u32 = 8; const SPATIAL_RADIUS: u32 = 4; @@ -36,6 +31,24 @@ struct Knobs { k_max: u32, sigma: f32, lambda_ht: f32, + mismatch_scale: f32, + /// Whether a temporal member's own match distance inflates its + /// noise variance. Every other test in this file relies on a + /// uniform `sigma^2` across the whole group, so this defaults off + /// and only the mismatch-variance test itself turns it on. + use_member_sigma: bool, + /// Half-width of each neighbour's refine window, defaulting to the + /// module's [`REFINE`]. + refine: u32, + /// The expected distance two noisy copies of the same content show + /// by chance, subtracted from a member's raw match distance before + /// it becomes mismatch variance. Every other test in this file + /// leaves this at `0.0`, so a member's raw distance passes through + /// unchanged. + noise_floor: f32, + /// The motion block side length, defaulting to the module's + /// [`BLKSIZE`]. At [`BLK_STEP`] exactly one block covers a patch. + blksize: u32, } impl Default for Knobs { @@ -45,6 +58,11 @@ impl Default for Knobs { k_max: K_MAX, sigma: 0.02, lambda_ht: 2.7, + mismatch_scale: 1.0, + use_member_sigma: false, + refine: REFINE, + noise_floor: 0.0, + blksize: BLKSIZE, } } } @@ -120,19 +138,19 @@ fn run_fused_over(fx: &RingFixture, k: Knobs) -> FusedRun { ArrayArg::from_raw_parts(wsum.clone(), pixels * frames), ArrayArg::from_raw_parts(group_weight.clone(), refs), fx.centre_slot, - 0.0f32, + k.noise_floor, k.c_min, - THSAD, + k.mismatch_scale * k.mismatch_scale, k.lambda_ht, weight_scale(k.sigma, &profile), cross_frame_accum_scale(SPATIAL_RADIUS, fx.radius), - false, + k.use_member_sigma, fx.radius, - REFINE, + k.refine, fx.mv_stride, fx.conf_stride, BLK_STEP, - BLKSIZE, + k.blksize, fx.blocks_x, fx.blocks_y, w, @@ -282,3 +300,338 @@ fn no_admission_gate_means_the_group_always_fills() { one-member run deposited" ); } + +/// A temporal member's extra variance is its own match distance, per +/// channel and per pixel, times the scale squared. +/// +/// `planted_ring` puts exact copies of the reference patch in every +/// neighbour. Adding a uniform offset `d` to each copy gives every +/// temporal member the distance `3 * 64 * d^2` and so the variance +/// `d^2 * scale^2`. With `lambda_ht` huge only the group DC survives, +/// whose variance is the ladder's level 0, and the group weight is its +/// reciprocal. `haar_variance_ladder` is the host mirror the GPU ladder +/// is already pinned against. +/// +/// The run uses `refine: 0`, which collapses each neighbour's window to +/// its single motion-predicted position, exactly where `planted_ring` +/// puts the copy. That makes the group composition exact — self, the +/// four planted copies, and three centre-frame spatial members with no +/// mismatch variance of their own — so the expected variance below can +/// be written down at all. A wider window admits near-miss candidates +/// that tie with genuine spatial ones and leak mismatch variance into +/// what should be a clean baseline. +#[test] +fn a_temporal_member_carries_its_own_match_distance_as_variance() { + use crate::collab::kernels::transforms::haar_variance_ladder; + + let (w, h) = (96u32, 96u32); + let radius = 2u32; + let ref_pos = (64u32, 64u32); + let patch = deterministic_texture(5); + let sigma = 0.02f32; + let refs_x = refs_along(w); + let ref_idx = ((ref_pos.1 / STEP) * refs_x + (ref_pos.0 / STEP)) as usize; + + for (d, scale) in [(0.0f32, 1.0f32), (0.05, 1.0), (0.05, 2.0), (0.1, 1.0)] { + let mut fx = planted_ring(w, h, radius, ref_pos, 3, &patch, 0.2, |_| 1.0); + // Offset every neighbour copy by d. The neighbour slots are + // every slot but the centre. + let pixels = (w * h) as usize; + for slot in 0..(2 * radius + 1) { + if slot == fx.centre_slot { + continue; + } + let frame = &mut fx.ring[slot as usize * pixels..(slot as usize + 1) * pixels]; + for v in frame.iter_mut() { + if *v > 0.5 { + *v += d; + } + } + } + + let run = run_fused_over( + &fx, + Knobs { + sigma, + lambda_ht: 1.0e6, + mismatch_scale: scale, + use_member_sigma: true, + refine: 0, + ..Knobs::default() + }, + ); + + // Members sort by distance: self, then the four temporal + // copies at 3 * 64 * d^2 each, then three flat spatial patches. + let base = sigma * sigma; + let mut v = [base; 8]; + for m in v.iter_mut().take(5).skip(1) { + *m = base + d * d * scale * scale; + } + let expected = 1.0 / haar_variance_ladder(&v, 8)[0]; + let got = run.group_weight[ref_idx]; + assert!( + (got - expected).abs() <= expected * 1e-3, + "d={d} scale={scale}: expected group weight {expected}, got {got}" + ); + } +} + +/// A non-zero `noise_floor` subtracts from a temporal member's raw +/// match distance before it becomes variance, so a larger floor lowers +/// the member's variance. +/// +/// Same fixture and offset as +/// [`a_temporal_member_carries_its_own_match_distance_as_variance`], at +/// `d = 0.1` and `scale = 1.0`, so each temporal member's raw distance +/// is `3 * 64 * d^2 = 1.92`. A floor of `0.96`, half that distance, +/// leaves excess `0.96` and so variance `0.96 / (3 * 64) = 0.005`, half +/// of the `d^2 = 0.01` a zero floor would give. +#[test] +fn a_noise_floor_lowers_a_temporal_members_variance_by_the_expected_amount() { + use crate::collab::kernels::transforms::haar_variance_ladder; + + let (w, h) = (96u32, 96u32); + let radius = 2u32; + let ref_pos = (64u32, 64u32); + let patch = deterministic_texture(5); + let sigma = 0.02f32; + let d = 0.1f32; + let refs_x = refs_along(w); + let ref_idx = ((ref_pos.1 / STEP) * refs_x + (ref_pos.0 / STEP)) as usize; + + let mut fx = planted_ring(w, h, radius, ref_pos, 3, &patch, 0.2, |_| 1.0); + let pixels = (w * h) as usize; + for slot in 0..(2 * radius + 1) { + if slot == fx.centre_slot { + continue; + } + let frame = &mut fx.ring[slot as usize * pixels..(slot as usize + 1) * pixels]; + for v in frame.iter_mut() { + if *v > 0.5 { + *v += d; + } + } + } + + let raw_distance = 3.0 * PATCH_AREA as f32 * d * d; + let noise_floor = raw_distance / 2.0; + + let run = run_fused_over( + &fx, + Knobs { + sigma, + lambda_ht: 1.0e6, + use_member_sigma: true, + refine: 0, + noise_floor, + ..Knobs::default() + }, + ); + + let base = sigma * sigma; + let excess = (raw_distance - noise_floor).max(0.0); + let member_variance = excess / (3.0 * PATCH_AREA as f32); + let mut v = [base; 8]; + for m in v.iter_mut().take(5).skip(1) { + *m = base + member_variance; + } + let expected = 1.0 / haar_variance_ladder(&v, 8)[0]; + let got = run.group_weight[ref_idx]; + assert!( + (got - expected).abs() <= expected * 1e-3, + "noise_floor={noise_floor}: expected group weight {expected} (member variance \ + {member_variance}), got {got}" + ); + + // The floor must actually have lowered the variance, not left it at + // the zero-floor value the previous test measured at this same d. + assert!( + member_variance < d * d, + "expected the floor to lower the member variance below the zero-floor value {}, got {}", + d * d, + member_variance + ); +} + +/// Sets one block's vector toward neighbour `t`. +fn set_block_mv(fx: &mut RingFixture, t: u32, bx: u32, by: u32, mv: [i32; 2]) { + let block = by * fx.blocks_x + bx; + let base = (t * fx.mv_stride + block * 2) as usize; + fx.mv_field[base] = mv[0]; + fx.mv_field[base + 1] = mv[1]; +} + +/// Writes an 8x8 patch into ring slot `slot` at `(px, py)`. +fn plant_in_slot(fx: &mut RingFixture, slot: u32, px: u32, py: u32, patch: &[f32; 64]) { + let pixels = (fx.width * fx.height) as usize; + let frame = &mut fx.ring[slot as usize * pixels..(slot as usize + 1) * pixels]; + for row in 0..8u32 { + for col in 0..8u32 { + frame[((py + row) * fx.width + px + col) as usize] = patch[(row * 8 + col) as usize]; + } + } +} + +/// Moves each neighbour's copy of the reference patch 20 pixels right, +/// leaving flat background where the reference sits, and points one +/// block's vector at the copy. +/// +/// `planted_ring` at a zero shift puts a copy at the reference position +/// in every frame, so the copy there is erased first. Every block but +/// `(bx, by)` then holds the zeroed vector `planted_ring` left, which +/// points at flat background, so the copy is reachable only through +/// `(bx, by)`. +fn only_reachable_through( + fx: &mut RingFixture, + ref_pos: (u32, u32), + patch: &[f32; 64], + (bx, by): (u32, u32), +) { + let flat = [0.2f32; 64]; + for t in 0..fx.neighbour_slots.len() as u32 { + let slot = fx.neighbour_slots[t as usize]; + plant_in_slot(fx, slot, ref_pos.0, ref_pos.1, &flat); + plant_in_slot(fx, slot, ref_pos.0 + 20, ref_pos.1, patch); + set_block_mv(fx, t, 8, 8, [0, 0]); + set_block_mv(fx, t, bx, by, [20, 0]); + } +} + +/// The corner block's vector points at flat background, and only a +/// neighbouring covering block's vector points at the planted copy. +/// +/// The reference at (64, 64) sits on the corner of block (8, 8) and is +/// also covered by blocks (7, 7), (8, 7) and (7, 8), since a 16-pixel +/// block at an 8-pixel step covers two patches per axis. A search that +/// reads only the corner block never sees the copy. +/// +/// Each of the three non-corner covering blocks is tried on its own, +/// `(7, 7)` diagonally, `(8, 7)` above and `(7, 8)` to the left, so a +/// kernel that read only the corner and the diagonal fails on two of +/// the three. +/// +/// The ring runs at radius 2, so four neighbours each hold a copy the +/// covering block reaches. Half the group is then an exact copy of the +/// reference, against a group of near-background patches when only the +/// corner block is read, and the group weight separates the two by a +/// wide margin. The control leaves every block on the corner's zeroed +/// vector, so no rectangle reaches the copy however many blocks are +/// read. +#[test] +fn a_covering_block_other_than_the_corner_finds_the_match() { + let (w, h) = (96u32, 96u32); + let radius = 2u32; + let ref_pos = (64u32, 64u32); + let patch = deterministic_texture(13); + let refs_x = refs_along(w); + let ref_idx = ((ref_pos.1 / STEP) * refs_x + (ref_pos.0 / STEP)) as usize; + + // The same ring with every vector zeroed, so no block's rectangle + // reaches the copy however many blocks are read. + let mut corner_only = planted_ring(w, h, radius, ref_pos, 0, &patch, 0.2, |_| 1.0); + only_reachable_through(&mut corner_only, ref_pos, &patch, (8, 8)); + corner_only.mv_field.fill(0); + let without = run_fused_over(&corner_only, Knobs::default()).group_weight[ref_idx]; + + for block in [(7u32, 7u32), (8, 7), (7, 8)] { + let mut fx = planted_ring(w, h, radius, ref_pos, 0, &patch, 0.2, |_| 1.0); + only_reachable_through(&mut fx, ref_pos, &patch, block); + let with_covering = run_fused_over(&fx, Knobs::default()).group_weight[ref_idx]; + + assert!( + with_covering > without * 1.5, + "the copies are only reachable through block {block:?}'s vector, expected a far \ + better group with it, got {with_covering} against {without}" + ); + } +} + +/// Two covering blocks whose vectors differ by one pixel give +/// overlapping rectangles, and a position inside both is scored once. +/// +/// A copy planted where both rectangles reach it would otherwise enter +/// the group twice. With `lambda_ht` huge every member deposits the +/// same weight, so the planted patch's pixels receive exactly the +/// weight they receive when only one block points at it. +#[test] +fn overlapping_covering_rectangles_score_each_position_once() { + let (w, h) = (96u32, 96u32); + let radius = 1u32; + let ref_pos = (64u32, 64u32); + let patch = deterministic_texture(17); + let pixels = (w * h) as usize; + let knobs = || Knobs { + lambda_ht: 1.0e6, + ..Knobs::default() + }; + + let build = |second_vector: Option<[i32; 2]>| { + let mut fx = planted_ring(w, h, radius, ref_pos, 0, &patch, 0.2, |_| 1.0); + for t in 0..2u32 { + let slot = fx.neighbour_slots[t as usize]; + plant_in_slot(&mut fx, slot, ref_pos.0 + 20, ref_pos.1, &patch); + set_block_mv(&mut fx, t, 8, 8, [20, 0]); + if let Some(v) = second_vector { + set_block_mv(&mut fx, t, 7, 7, v); + } + } + fx + }; + + let one = run_fused_over(&build(None), knobs()); + let two = run_fused_over(&build(Some([21, 0])), knobs()); + + let frames = one.wsum.len() / pixels; + let planted_centre = + |run: &FusedRun, s: usize| run.wsum[s * pixels + ((ref_pos.1 + 4) * w + ref_pos.0 + 20 + 4) as usize]; + for s in 0..frames { + if s as u32 == 1 { + continue; + } + assert_eq!( + planted_centre(&two, s), + planted_centre(&one, s), + "slot {s}: the planted copy must carry the same weight whether one or two covering \ + blocks reach it" + ); + } + assert!( + planted_centre(&one, 0) > 0, + "the copy must be a member in the first place" + ); +} + +/// With `blksize == step` exactly one block covers a patch, so a +/// neighbouring block's vector is never consulted. +/// +/// The copy is reachable only through block `(7, 7)`, which covers the +/// patch at `blksize = 16` and does not at `blksize = 8`. +#[test] +fn a_block_size_equal_to_the_step_reads_only_the_corner_block() { + let (w, h) = (96u32, 96u32); + let radius = 2u32; + let ref_pos = (64u32, 64u32); + let patch = deterministic_texture(19); + let refs_x = refs_along(w); + let ref_idx = ((ref_pos.1 / STEP) * refs_x + (ref_pos.0 / STEP)) as usize; + + let mut fx = planted_ring(w, h, radius, ref_pos, 0, &patch, 0.2, |_| 1.0); + only_reachable_through(&mut fx, ref_pos, &patch, (7, 7)); + + let covering = run_fused_over(&fx, Knobs::default()).group_weight[ref_idx]; + let single = run_fused_over( + &fx, + Knobs { + blksize: BLK_STEP, + ..Knobs::default() + }, + ) + .group_weight[ref_idx]; + + assert!( + covering > single * 1.5, + "at blksize == step the copy is unreachable, got {single} against {covering} with \ + covering blocks" + ); +} diff --git a/av-denoise-core/src/nl4d/tests/mod.rs b/av-denoise-core/src/nl4d/tests/mod.rs index c47cd8b..689d12b 100644 --- a/av-denoise-core/src/nl4d/tests/mod.rs +++ b/av-denoise-core/src/nl4d/tests/mod.rs @@ -3,3 +3,4 @@ mod helpers; mod confidence; mod grouping; mod pipeline; +mod regularise; diff --git a/av-denoise-core/src/nl4d/tests/pipeline.rs b/av-denoise-core/src/nl4d/tests/pipeline.rs index b2fd0f8..72015a4 100644 --- a/av-denoise-core/src/nl4d/tests/pipeline.rs +++ b/av-denoise-core/src/nl4d/tests/pipeline.rs @@ -58,6 +58,7 @@ fn static_clip_params(temporal_radius: u32) -> Nl4dParams { // The shipped default, so these run the aggregation a real // caller gets. kaiser_beta: 2.0, + field_lambda: 0.0, } } @@ -703,7 +704,9 @@ fn cross_frame_aggregation_beats_centre_only_at_the_same_lambda() { centre_slot, 0.0f32, C_MIN, - front.thsad_value(), + // `use_member_sigma` is off below, so this never reaches + // a threshold and any value is exact. + 1.0f32, LAMBDA_HT, wnorm, accum_scale, @@ -770,3 +773,223 @@ fn cross_frame_aggregation_beats_centre_only_at_the_same_lambda() { lambda_ht, got cross-frame={cross_frame_psnr:.4} dB centre-only={centre_only_psnr:.4} dB" ); } + +/// The snapshot reports the field the fused kernel was given. A clip +/// whose every frame is the previous one shifted right by 2 pixels must +/// report `[2 * k, 0]` toward the neighbour at offset `k`, at an +/// interior block, once the first pass has run. +#[test] +fn motion_snapshot_reports_the_field_the_pass_used() { + let client = make_client(); + let (w, h) = (96u32, 96u32); + let radius = 2u32; + let base = textured_base(w, h); + let frames: Vec> = (0..5i32) + .map(|k| { + let mut f = vec![0.0f32; (w * h) as usize]; + for y in 0..h { + for x in 0..w { + let sx = (x as i32 - 2 * (k - 2)).clamp(0, w as i32 - 1) as u32; + f[(y * w + x) as usize] = base[(y * w + sx) as usize]; + } + } + f + }) + .collect(); + + let mut d = + Nl4dDenoiser::::new(&client, static_clip_params(radius), w, h).expect("construction failed"); + assert!(d.motion_snapshot().is_none(), "no pass has run yet"); + for frame in &frames { + d.push_frame(frame); + let _ = d.denoise_submit().expect("denoise_submit failed"); + } + + let snap = d.motion_snapshot().expect("a pass has run"); + assert_eq!(snap.offsets, vec![-2, -1, 1, 2]); + assert_eq!(snap.step, 8); + assert_eq!(snap.blksize, 16); + // Block (3, 3) covers pixels 24..40, well inside the frame. + let block = (3 * snap.blocks_x + 3) as usize; + for (t, &k) in snap.offsets.iter().enumerate() { + assert_eq!( + snap.vectors[t][block], + [2 * k, 0], + "neighbour k={k} should be tracked as a 2*k pixel shift" + ); + assert!( + snap.confidence[t][block] > 0.5, + "a clean shift must score confidently" + ); + } +} + +/// A panning clip with one flat block. The estimator ties on the flat +/// block and leaves it at the seed, so its vector differs from its +/// neighbours'. With `field_lambda` on, the pass pulls it to the +/// neighbourhood's vector, and the snapshot shows the regularised +/// field. With it off the field is the estimator's. +#[test] +fn field_regularisation_reaches_the_snapshot() { + let client = make_client(); + let (w, h) = (128u32, 96u32); + let radius = 1u32; + let mut base = textured_base(w, h); + // Flatten a 24x24 region centred on block (7, 5)'s own footprint, + // 52..76 x 36..60. Large enough to tie both the fine-level search and + // the coarse pyramid level for that one block, but small enough that + // its overlapping neighbours still see enough texture past the + // region's edge to estimate correctly, so only the centre block ties. + for y in 36..60u32 { + for x in 52..76u32 { + base[(y * w + x) as usize] = 0.5; + } + } + let frames: Vec> = (0..3i32) + .map(|k| { + let mut f = vec![0.0f32; (w * h) as usize]; + for y in 0..h { + for x in 0..w { + let sx = (x as i32 - 3 * (k - 1)).clamp(0, w as i32 - 1) as u32; + f[(y * w + x) as usize] = base[(y * w + sx) as usize]; + } + } + f + }) + .collect(); + + let run = |lambda: f32| { + let params = Nl4dParams { + field_lambda: lambda, + ..static_clip_params(radius) + }; + let mut d = Nl4dDenoiser::::new(&client, params, w, h).expect("construction failed"); + for frame in &frames { + d.push_frame(frame); + let _ = d.denoise_submit().expect("denoise_submit failed"); + } + d.motion_snapshot().expect("a pass ran") + }; + + let off = run(0.0); + let on = run(1.0); + // The block at (7, 5) spans pixels 56..72 x 40..56, inside the flat + // region on every frame. + let flat_block = (5 * off.blocks_x + 7) as usize; + let t_plus = 1usize; + assert_ne!( + off.vectors[t_plus][flat_block], + [3, 0], + "the flat block must not be tracked without help, or this test proves nothing" + ); + assert_eq!(on.vectors[t_plus][flat_block], [3, 0]); + // A textured block is unchanged by the pass. + let textured_block = (2 * off.blocks_x + 2) as usize; + assert_eq!(off.vectors[t_plus][textured_block], [3, 0]); + assert_eq!(on.vectors[t_plus][textured_block], [3, 0]); +} + +/// `Nl4dParams::default()`, changing only `channels`, which has to +/// switch to `Luma` because this file's helpers only ever synthesise a +/// single plane. Every other test in this file pins `field_lambda: +/// 0.0` so its recorded values stay stable, but the shipped default is +/// `1.0`, meaning a real caller always runs the field-regularisation +/// pass this configuration exercises. +fn shipped_default_params() -> Nl4dParams { + let params = Nl4dParams { + nlm: NlmParams { + channels: ChannelMode::Luma, + ..Nl4dParams::default().nlm + }, + ..Nl4dParams::default() + }; + assert_eq!( + params.field_lambda, 1.0, + "this helper exists to exercise the shipped default, not an override" + ); + params +} + +/// Runs the pipeline at the true shipped defaults, in two phases. +/// +/// The first phase pushes a static, clean (noiseless) clip and checks +/// the field-regularisation pass leaves an interior block's vector at +/// exactly zero. A static clip's true motion is zero everywhere, so a +/// correctly regularised field has nothing to pull a well-textured +/// interior block's vector away from zero toward: every neighbouring +/// block's vector is also zero, so their median is zero, which is +/// already where the block sits. A field-regularisation dispatch that +/// reads the wrong pyramid slot, indexes the wrong neighbour, or +/// strides into a different block's data instead of its own would +/// instead pull in whatever mismatched vector sits there, and a real +/// motion vector would show up in a scene where nothing ever moved. +/// This phase carries no noise, because the estimator's own +/// noise-driven wobble would otherwise mask exactly the kind of small, +/// wrong-source displacement it exists to catch. +/// +/// The second phase pushes a static, noisy clip through a fresh +/// denoiser at the same defaults and checks every emitted frame comes +/// out well above the noisy input's own PSNR, the same property +/// [`denoises_a_static_noisy_clip`] checks at `field_lambda: 0.0`. A +/// field-regularisation dispatch bug severe enough to corrupt the +/// motion field would feed the temporal grouping kernel the wrong +/// candidates and show up here as a smaller improvement, or none at +/// all. +#[test] +fn shipped_defaults_denoise_a_static_clip_and_regularise_its_field_to_zero() { + let client = make_client(); + let (w, h) = (96u32, 96u32); + let base = textured_base(w, h); + let radius = shipped_default_params().temporal_radius; + let n = (3 * radius + 1) as usize; + + // Phase 1: a clean, static clip, checking the regularised field. + let mut clean_d = Nl4dDenoiser::::new(&client, shipped_default_params(), w, h) + .expect("construction failed for the clean phase"); + for _ in 0..n { + clean_d.push_frame(&base); + let _ = clean_d.denoise_submit().expect("denoise_submit failed"); + } + let snap = clean_d.motion_snapshot().expect("a pass ran"); + // Block (2, 2) spans pixels 16..32 on both axes, well inside the + // frame and away from any edge-clamping effects. + let interior_block = (2 * snap.blocks_x + 2) as usize; + for (t, &k) in snap.offsets.iter().enumerate() { + assert_eq!( + snap.vectors[t][interior_block], + [0, 0], + "neighbour k={k}: a static, noiseless clip's regularised field must read exactly \ + zero at an interior block; a nonzero vector here is what a wrong pyramid slot, \ + neighbour index, or stride in the smoothing dispatch looks like" + ); + } + + // Phase 2: a noisy version of the same clip, checking the output. + let frames: Vec> = (0..n as u32) + .map(|seed| noisy_copy_of(&base, w, h, SIGMA, seed)) + .collect(); + let mut noisy_d = Nl4dDenoiser::::new(&client, shipped_default_params(), w, h) + .expect("construction failed for the noisy phase"); + let mut outputs: Vec> = Vec::new(); + for frame in &frames { + noisy_d.push_frame(frame); + if let Some(pending) = noisy_d.denoise_submit().expect("denoise_submit failed") { + let frame = pending.wait().expect("readback failed"); + outputs.push(frame.into_f32().expect("f32 output")); + } + } + noisy_d + .flush(|frame| outputs.push(frame.as_f32().expect("f32 denoiser").to_vec())) + .expect("flush failed"); + + assert_eq!(outputs.len(), n, "expected one emitted frame per pushed frame"); + for (i, out) in outputs.iter().enumerate() { + let noisy_psnr = psnr(&frames[i], &base); + let out_psnr = psnr(out, &base); + assert!( + out_psnr > noisy_psnr + 6.0, + "frame {i}: expected at least a 6 dB PSNR improvement over the noisy input at the \ + shipped defaults, got noisy={noisy_psnr:.4} dB denoised={out_psnr:.4} dB" + ); + } +} diff --git a/av-denoise-core/src/nl4d/tests/regularise.rs b/av-denoise-core/src/nl4d/tests/regularise.rs new file mode 100644 index 0000000..8d9aafe --- /dev/null +++ b/av-denoise-core/src/nl4d/tests/regularise.rs @@ -0,0 +1,255 @@ +use cubecl::prelude::*; + +use super::helpers::{R, make_client}; +use crate::nl4d::kernels::nl4d_mv_regularise; +use crate::nlmeans::motion::THSAD_PIXEL; + +const BLKSIZE: u32 = 16; +const STEP: u32 = 8; + +/// One launch over a `blocks_x x blocks_y` grid, returning the output +/// field and confidence. +fn run( + w: u32, + h: u32, + centre: &[f32], + neighbour: &[f32], + mv_in: &[i32], + lambda: f32, +) -> (Vec, Vec) { + let client = make_client(); + let blocks_x = w.div_ceil(STEP); + let blocks_y = h.div_ceil(STEP); + let blocks = (blocks_x * blocks_y) as usize; + assert_eq!(mv_in.len(), 2 * blocks); + let centre_buf = client.create_from_slice(f32::as_bytes(centre)); + let neighbour_buf = client.create_from_slice(f32::as_bytes(neighbour)); + let mv_in_buf = client.create_from_slice(i32::as_bytes(mv_in)); + let mv_out = client.empty(2 * blocks * size_of::()); + let conf_out = client.empty(blocks * size_of::()); + let thsad = (BLKSIZE * BLKSIZE) as f32 * THSAD_PIXEL; + + unsafe { + nl4d_mv_regularise::launch_unchecked::( + &client, + CubeCount::new_2d(blocks_x, blocks_y), + CubeDim::new_2d(8, 8), + ArrayArg::from_raw_parts(centre_buf, centre.len()), + ArrayArg::from_raw_parts(neighbour_buf, neighbour.len()), + ArrayArg::from_raw_parts(mv_in_buf, 2 * blocks), + ArrayArg::from_raw_parts(mv_out.clone(), 2 * blocks), + ArrayArg::from_raw_parts(conf_out.clone(), blocks), + lambda * (BLKSIZE * BLKSIZE) as f32 * THSAD_PIXEL, + 0.0, + thsad, + w, + h, + BLKSIZE, + STEP, + blocks_x, + blocks_y, + ); + } + + let mv = i32::from_bytes(&client.read_one(mv_out).expect("mv readback"))[..2 * blocks].to_vec(); + let conf = f32::from_bytes(&client.read_one(conf_out).expect("conf readback"))[..blocks].to_vec(); + (mv, conf) +} + +/// A frame with distinct values everywhere. +fn textured(w: u32, h: u32, seed: u32) -> Vec { + (0..w * h) + .map(|i| { + let mut x = i + .wrapping_mul(2654435761) + .wrapping_add(seed.wrapping_mul(0x9E37_79B9)); + x ^= x >> 15; + x = x.wrapping_mul(0x85EB_CA6B); + x ^= x >> 13; + 0.2 + 0.6 * (x as f32 / u32::MAX as f32) + }) + .collect() +} + +/// `neighbour(x, y) = centre(x - dx, y - dy)`, so the true vector is +/// `(dx, dy)`. +fn shifted(centre: &[f32], w: u32, h: u32, dx: i32, dy: i32) -> Vec { + let mut out = vec![0.0f32; (w * h) as usize]; + for y in 0..h as i32 { + for x in 0..w as i32 { + let sx = (x - dx).clamp(0, w as i32 - 1) as u32; + let sy = (y - dy).clamp(0, h as i32 - 1) as u32; + out[(y as u32 * w + x as u32) as usize] = centre[(sy * w + sx) as usize]; + } + } + out +} + +fn uniform_field(blocks: usize, v: [i32; 2]) -> Vec { + let mut f = Vec::with_capacity(2 * blocks); + for _ in 0..blocks { + f.push(v[0]); + f.push(v[1]); + } + f +} + +/// A flat centre and neighbour score the same SAD at every candidate, +/// so an outlier vector in a smooth field moves to the neighbourhood's +/// median as soon as the penalty is positive. +#[test] +fn an_outlier_in_a_flat_region_moves_to_the_median() { + let (w, h) = (64u32, 64u32); + let blocks_x = w.div_ceil(STEP); + let blocks = (blocks_x * h.div_ceil(STEP)) as usize; + let flat = vec![0.5f32; (w * h) as usize]; + let mut field = uniform_field(blocks, [3, 1]); + let outlier = (4 * blocks_x + 4) as usize; + field[2 * outlier] = -6; + field[2 * outlier + 1] = 5; + + let (out, _) = run(w, h, &flat, &flat, &field, 1.0); + assert_eq!([out[2 * outlier], out[2 * outlier + 1]], [3, 1]); + // Every other block already sits on its median and stays put. + for b in 0..blocks { + if b != outlier { + assert_eq!([out[2 * b], out[2 * b + 1]], [3, 1], "block {b} moved"); + } + } +} + +/// `lambda = 0` is a plain re-score. On a clean shift with the field +/// already correct nothing moves, and on a flat region ties go to the +/// block's own vector, so the outlier stays. +#[test] +fn a_zero_penalty_keeps_the_input_field_on_ties() { + let (w, h) = (64u32, 64u32); + let blocks_x = w.div_ceil(STEP); + let blocks = (blocks_x * h.div_ceil(STEP)) as usize; + let flat = vec![0.5f32; (w * h) as usize]; + let mut field = uniform_field(blocks, [3, 1]); + let outlier = (4 * blocks_x + 4) as usize; + field[2 * outlier] = -6; + field[2 * outlier + 1] = 5; + + let (out, _) = run(w, h, &flat, &flat, &field, 0.0); + assert_eq!(out, field); +} + +/// A block whose own vector matches far better than the median keeps +/// it, because its SAD margin exceeds the penalty. +#[test] +fn a_true_boundary_block_keeps_its_vector_when_the_sad_margin_wins() { + let (w, h) = (64u32, 64u32); + let blocks_x = w.div_ceil(STEP); + let blocks = (blocks_x * h.div_ceil(STEP)) as usize; + let centre = textured(w, h, 1); + // The whole neighbour is the centre shifted by (2, 0). + let neighbour = shifted(¢re, w, h, 2, 0); + // The field says (0, 0) everywhere except one interior block that + // knows the truth. + let mut field = uniform_field(blocks, [0, 0]); + let truthful = (4 * blocks_x + 4) as usize; + field[2 * truthful] = 2; + + let (out, conf) = run(w, h, ¢re, &neighbour, &field, 1.0); + assert_eq!([out[2 * truthful], out[2 * truthful + 1]], [2, 0]); + assert!( + conf[truthful] > 0.9, + "an exact match scores a high confidence, got {}", + conf[truthful] + ); + // Its neighbours see (2, 0) among the adjacent candidates and take + // it, since textured content beats the penalty of one pixel. + let right = truthful + 1; + assert_eq!([out[2 * right], out[2 * right + 1]], [2, 0]); +} + +/// The median rule picks the lower of two middle values, not the upper, +/// and a corner block's three-member neighbourhood gets exercised with +/// genuinely different values along the way. +/// +/// The field is flat, so every candidate scores the same zero SAD and +/// only the penalty against the median decides the winner. The centre +/// block's eight neighbours split four and four between `-1` and `2`, +/// so the lower median is `-1` while the upper median would be `2`, and +/// the kernel must land on `-1`. Block `(0, 0)` is a grid corner with +/// only three neighbours, one of which is the centre block, so its +/// median comes from the mismatched values `-1, -1, 99` rather than +/// eight identical entries. +#[test] +fn the_median_rule_picks_the_lower_of_two_middle_values() { + let (w, h) = (24u32, 24u32); + let blocks_x = w.div_ceil(STEP); + let blocks_y = h.div_ceil(STEP); + assert_eq!((blocks_x, blocks_y), (3, 3)); + let flat = vec![0.5f32; (w * h) as usize]; + + #[rustfmt::skip] + let field: Vec = vec![ + -1, 0, -1, 0, -1, 0, + -1, 0, 99, 99, 2, 0, + 2, 0, 2, 0, 2, 0, + ]; + + let (out, _) = run(w, h, &flat, &flat, &field, 1.0); + let centre = (blocks_x + 1) as usize; + assert_eq!( + [out[2 * centre], out[2 * centre + 1]], + [-1, 0], + "the centre block must take the lower median, not the upper one" + ); + let corner = 0usize; + assert_eq!( + [out[2 * corner], out[2 * corner + 1]], + [-1, 0], + "the corner block's three-member median must resolve too" + ); +} + +/// Confidence is recomputed for the winner, not copied from the input. +#[test] +fn confidence_follows_the_winning_vector() { + let (w, h) = (64u32, 64u32); + let blocks_x = w.div_ceil(STEP); + let blocks = (blocks_x * h.div_ceil(STEP)) as usize; + let centre = textured(w, h, 2); + let neighbour = shifted(¢re, w, h, 1, 1); + let field = uniform_field(blocks, [1, 1]); + + let (_, conf) = run(w, h, ¢re, &neighbour, &field, 1.0); + let interior = (3 * blocks_x + 3) as usize; + assert!( + conf[interior] > 0.99, + "a perfect match must score ~1, got {}", + conf[interior] + ); + + let wrong = uniform_field(blocks, [-3, -3]); + let (_, conf) = run(w, h, ¢re, &neighbour, &wrong, 0.0); + assert!( + conf[interior] < 0.5, + "a wrong vector on texture must score low, got {}", + conf[interior] + ); + + // A block whose winner is not its own vector, so this test does not + // pass merely by reporting candidate 0's confidence unconditionally. + // `right` sits beside a block that knows the true shift, and its own + // vector is wrong, so it wins on its left neighbour's vector, which is + // candidate 2. + let boundary_centre = textured(w, h, 1); + let boundary_neighbour = shifted(&boundary_centre, w, h, 2, 0); + let mut boundary_field = uniform_field(blocks, [0, 0]); + let truthful = (4 * blocks_x + 4) as usize; + boundary_field[2 * truthful] = 2; + let right = truthful + 1; + + let (out, conf) = run(w, h, &boundary_centre, &boundary_neighbour, &boundary_field, 1.0); + assert_eq!([out[2 * right], out[2 * right + 1]], [2, 0]); + assert!( + conf[right] > 0.9, + "right wins its neighbour's exact match (candidate 2), so confidence must be high, got {}", + conf[right] + ); +} diff --git a/av-denoise-core/src/nlmeans/denoiser.rs b/av-denoise-core/src/nlmeans/denoiser.rs index 7e78a87..1623cf0 100644 --- a/av-denoise-core/src/nlmeans/denoiser.rs +++ b/av-denoise-core/src/nlmeans/denoiser.rs @@ -80,6 +80,13 @@ pub(crate) struct RingView { pub mv_stride: u32, /// `f32` element stride between neighbours in `confidence`. pub conf_stride: u32, + /// The luma pyramid the motion estimator analysed, the reference + /// ring's when a prefilter is active and the input ring's + /// otherwise. Level 0 of slot `s` starts at + /// `pyramid_slot_byte_offset(width, height, frame_count, 0, s, align)`. + pub pyramid: Handle, + /// How many frames the ring holds. + pub frame_count: u32, } /// The stateful NLMeans denoiser that owns the GPU buffers. @@ -1599,6 +1606,12 @@ impl NlmDenoiser { )) })? .clone(); + let pyramid = self + .pyramid_reference + .as_ref() + .or(self.pyramid_input.as_ref()) + .expect("pyramid allocated when mc_ctx is Some") + .clone(); Ok(Some(RingView { input: self.input_buf.clone(), @@ -1608,6 +1621,8 @@ impl NlmDenoiser { neighbour_slots, mv_stride: (mc.mv_field_bytes_per_neighbour() / size_of::() as u64) as u32, conf_stride: (mc.confidence_bytes_per_neighbour() / size_of::() as u64) as u32, + pyramid, + frame_count: self.params.total_frames(), })) } @@ -1643,8 +1658,21 @@ impl NlmDenoiser { .expect("motion_ctx called without motion compensation active") } - /// `thsad(blksize, thsad_scale)` in normalised SAD units, the same - /// threshold [`Self::submit_machinery`] scores confidence against. + /// The SAD two noisy copies of one block show by chance, the floor + /// [`Self::submit_machinery`] scored confidence against. + /// + /// # Panics + /// + /// Panics under the same condition as [`Self::motion_ctx`]. + pub(crate) fn sad_noise_floor_value(&self) -> f32 { + let blksize = self.motion_ctx().blksize; + let sigma = crate::nlmeans::dispatch::mc_sad_noise_floor_sigma(self.params.prefilter, self.sigma_y); + motion::sad_noise_floor(blksize, sigma) + } + + /// The SAD threshold [`Self::submit_machinery`] scored confidence + /// against, the same one [`Self::sad_noise_floor_value`] is measured + /// past. /// /// # Panics /// diff --git a/av-denoise-core/src/nlmeans/dispatch.rs b/av-denoise-core/src/nlmeans/dispatch.rs index 8b868bb..d23d569 100644 --- a/av-denoise-core/src/nlmeans/dispatch.rs +++ b/av-denoise-core/src/nlmeans/dispatch.rs @@ -138,7 +138,7 @@ const BILATERAL_RESIDUAL_FRACTION: f32 = 0.0; /// An `External` reference comes from the caller with unknown noise, and /// is not something this crate denoised, so it keeps the raw sigma just /// as `PrefilterMode::None` does. -fn mc_sad_noise_floor_sigma(prefilter: PrefilterMode, sigma_y: f32) -> f32 { +pub(super) fn mc_sad_noise_floor_sigma(prefilter: PrefilterMode, sigma_y: f32) -> f32 { match prefilter { PrefilterMode::NlmSpatial { .. } => sigma_y * NLM_SPATIAL_RESIDUAL_FRACTION, PrefilterMode::Bilateral { .. } => sigma_y * BILATERAL_RESIDUAL_FRACTION, diff --git a/av-denoise-core/src/nlmeans/kernels/helpers.rs b/av-denoise-core/src/nlmeans/kernels/helpers.rs index 956a87f..88006cb 100644 --- a/av-denoise-core/src/nlmeans/kernels/helpers.rs +++ b/av-denoise-core/src/nlmeans/kernels/helpers.rs @@ -78,6 +78,66 @@ pub(crate) fn channel_scale(#[comptime] channels: u32) -> f32 { scale } +/// The host mirror of [`channel_scale`]. +pub fn channel_scale_host(channels: u32) -> f32 { + match channels { + 1 => 3.0, + 2 => 1.5, + _ => 1.0, + } +} + +#[cfg(all(test, any(feature = "vulkan", feature = "metal")))] +mod tests { + use cubecl::prelude::*; + use cubecl::wgpu::WgpuRuntime; + + use super::{channel_scale, channel_scale_host}; + + type R = WgpuRuntime; + + fn make_client() -> ComputeClient { + let device = ::Device::default(); + R::client(&device) + } + + /// Runs [`channel_scale`] on the GPU for all three channel counts it + /// ever runs with, so the host mirror is checked against the code + /// the filter actually calls rather than against itself. + #[cube(launch_unchecked)] + fn channel_scale_probe(out: &mut Array) { + out[0] = channel_scale(1u32); + out[1] = channel_scale(2u32); + out[2] = channel_scale(3u32); + } + + #[test] + fn channel_scale_host_matches_the_kernel_mirror() { + let client = make_client(); + let out_buf = client.empty(3 * size_of::()); + + unsafe { + channel_scale_probe::launch_unchecked::( + &client, + CubeCount::new_1d(1), + CubeDim::new_1d(1), + ArrayArg::from_raw_parts(out_buf.clone(), 3), + ); + } + + let bytes = client.read_one(out_buf).expect("channel_scale readback failed"); + let got = f32::from_bytes(&bytes)[..3].to_vec(); + + for (idx, channels) in [1u32, 2, 3].into_iter().enumerate() { + assert_eq!( + got[idx], + channel_scale_host(channels), + "channels={channels}: host mirror disagrees with the GPU kernel" + ); + } + } +} + /// The Welsch weight for a box-summed patch distance. /// /// `noise_offset` is the distance two noisy copies of the same content diff --git a/av-denoise-core/src/nlmeans/motion/mod.rs b/av-denoise-core/src/nlmeans/motion/mod.rs index 851f3a5..2f02888 100644 --- a/av-denoise-core/src/nlmeans/motion/mod.rs +++ b/av-denoise-core/src/nlmeans/motion/mod.rs @@ -30,17 +30,15 @@ mod compensate; mod confidence; mod pyramid; +pub(crate) use analyse::{confidence_byte_offset, mv_field_byte_offset, run_analyse, run_seeded_refine}; #[cfg(all(test, any(feature = "vulkan", feature = "metal")))] -pub(crate) use analyse::mv_field_byte_offset; -pub(crate) use analyse::{confidence_byte_offset, run_analyse, run_seeded_refine}; -#[cfg(all(test, any(feature = "vulkan", feature = "metal")))] -pub(crate) use chain::{neighbour_idx_for_k, pair_byte_offset}; -pub(crate) use chain::{run_pair_analyse, zero_pair_slot}; +pub(crate) use chain::pair_byte_offset; +pub(crate) use chain::{neighbour_idx_for_k, run_pair_analyse, zero_pair_slot}; pub(crate) use compensate::run_compensate; -pub(crate) use confidence::{run_confidence_for_neighbour, sad_noise_floor, thsad}; +pub(crate) use confidence::{THSAD_PIXEL, run_confidence_for_neighbour, sad_noise_floor, thsad}; use cubecl::prelude::*; use cubecl::server::Handle; -pub(crate) use pyramid::{pyramid_pixels_per_frame, run_pyramid_build}; +pub(crate) use pyramid::{level_dims, pyramid_pixels_per_frame, pyramid_slot_byte_offset, run_pyramid_build}; use crate::nlmeans::align::StorageAlign; diff --git a/av-denoise-core/src/nlmeans/tests/machinery.rs b/av-denoise-core/src/nlmeans/tests/machinery.rs index e752570..6a976d8 100644 --- a/av-denoise-core/src/nlmeans/tests/machinery.rs +++ b/av-denoise-core/src/nlmeans/tests/machinery.rs @@ -151,6 +151,30 @@ fn submit_machinery_reports_ring_view_with_correct_motion_and_confidence() { ); } +/// The ring view carries the pyramid the estimator analysed and the +/// window size, and the front end reports the SAD noise floor it scored +/// confidence with. +#[test] +fn ring_view_exposes_the_analysed_pyramid_and_the_noise_floor() { + let client = make_client(); + let mut d = push_translating_sequence(&client); + let view = d + .submit_machinery() + .expect("submit_machinery dispatch failed") + .expect("window is exactly full, submit_machinery should report Some"); + + let frames = machinery_params().total_frames(); + assert_eq!(view.frame_count, frames); + // The pyramid holds every level of every slot, so it is at least + // one full-resolution luma plane per slot. + let bytes = client + .read_one(view.pyramid.clone()) + .expect("pyramid readback failed"); + let plane = bytes.len() / (frames as usize * size_of::()); + assert!(plane > 0, "the pyramid must hold at least one plane per slot"); + assert!(d.sad_noise_floor_value() >= 0.0); +} + /// A window that has not filled yet reports `None`, the same convention /// `denoise_submit_gpu` uses. #[test] diff --git a/av-denoise-vs/README.md b/av-denoise-vs/README.md index aeaca08..1a87cd8 100644 --- a/av-denoise-vs/README.md +++ b/av-denoise-vs/README.md @@ -95,8 +95,9 @@ clean.set_output() `lambda_ht_scale` is the threshold multiplier a transform coefficient's estimated-noise standard deviations must clear to survive. Raising it removes more noise and takes more fine detail with it. Try it in steps of about 0.05 before reaching for `lambda_ht`, -which pins luma and chroma's thresholds (5.3 and 4.2 by default) to the same absolute -number and loses that separation. +which pins luma and chroma's thresholds (5.2 and 3.4 by default, these values have been +manually tuned to provide the subjectively best image for a given grain strength across +real clips rather than synthetic benchmarks). `spatial_radius` is the speed dial. `preset` already resolves it, so setting `spatial_radius` explicitly overrides whatever the preset would have picked. The diff --git a/av-denoise/src/bin/cli/nl4d.rs b/av-denoise/src/bin/cli/nl4d.rs index a6d8192..3b10dcd 100644 --- a/av-denoise/src/bin/cli/nl4d.rs +++ b/av-denoise/src/bin/cli/nl4d.rs @@ -99,16 +99,17 @@ pub struct Nl4dArgs { /// How much a poorly matched neighbour patch is distrusted. /// - /// A neighbour block that motion tracking matched badly is treated - /// as a noisier view of the same content, and this scales how much - /// noisier. `1.0` (the library default) is the shipped calibration. - /// `0` matches `--no-confidence-variance`. + /// A patch matched in a neighbour frame is treated as a noisier view + /// of the same content, as noisy as its own match residual says, and + /// this scales how much noisier. `1.0` (the library default) is the + /// shipped calibration. `0` matches `--no-confidence-variance`. /// /// The variance grows with the square of this, so `2` distrusts a - /// bad match four times as much. The effect saturates somewhere - /// between `3` and `13` depending on how noisy the source is, and - /// values above `16` are rejected because nothing up there can - /// change a pixel. + /// bad match four times as much. The effect saturates. It saturates + /// sooner the worse the patch matched, because the variance the + /// mechanism derives is capped at 64 times the channel's own + /// variance, and values above `16` are rejected because nothing up + /// there can change a pixel. /// /// Setting this applies one value to both planes, unless /// `--luma-mismatch-scale` or `--chroma-mismatch-scale` overrides @@ -188,12 +189,21 @@ pub struct Nl4dArgs { #[arg(long)] pub kaiser_beta: Option, + /// How strongly motion vectors are pulled toward their neighbours. + /// + /// `1.0` (the library default) is the shipped calibration and + /// smooths the tracked field. Raise it to smooth out stray vectors + /// on noisy or flat content, at the cost of following small + /// objects less closely. `0` leaves the tracked field as it is. + #[arg(long)] + pub field_lambda: Option, + /// Stops a poorly matched patch from being trusted less than a well /// matched one. /// - /// On by default, the shrinkage treats a patch matched across frames - /// with low motion confidence as a noisier observation. This flag - /// gives every patch the same noise estimate instead. + /// On by default, the shrinkage treats a patch matched with a large + /// match distance as a noisier observation. This flag gives every + /// patch the same noise estimate instead. #[arg(long)] pub no_confidence_variance: bool, @@ -286,6 +296,7 @@ impl Nl4dArgs { mismatch_scale: self.mismatch_scale.unwrap_or(defaults.mismatch_scale), confidence_variance: !self.no_confidence_variance, kaiser_beta: self.kaiser_beta.unwrap_or(defaults.kaiser_beta), + field_lambda: self.field_lambda.unwrap_or(defaults.field_lambda), // The CLI keeps the temporal EMA every calibrated // preset assumes by default. Only `av-denoise-vs` // needs window-local estimation, for random-access @@ -454,6 +465,25 @@ mod tests { assert!((expect_nl4d(&opts).lambda_ht_scale - defaults.lambda_ht_scale).abs() < f32::EPSILON); } + #[test] + fn field_lambda_flows_into_the_nl4d_algorithm() { + let (args, nl4d) = parse(&["--field-lambda", "0.7"]); + let opts = nl4d.build_options(&args).expect("build_options should succeed"); + + assert!((nl4d.field_lambda.unwrap() - 0.7).abs() < f32::EPSILON); + assert!((expect_nl4d(&opts).field_lambda - 0.7).abs() < f32::EPSILON); + } + + #[test] + fn unset_field_lambda_resolves_to_the_library_default() { + let (args, nl4d) = parse(&[]); + let opts = nl4d.build_options(&args).expect("build_options should succeed"); + let defaults = Nl4dOptions::default(); + + assert_eq!(nl4d.field_lambda, None); + assert!((expect_nl4d(&opts).field_lambda - defaults.field_lambda).abs() < f32::EPSILON); + } + /// nl4d never runs an NLM weighting pass, so the flags that only /// configure one are gone rather than silently ignored. #[test] diff --git a/docs/TUNING-CLI.md b/docs/TUNING-CLI.md index f3d9948..e296a2b 100644 --- a/docs/TUNING-CLI.md +++ b/docs/TUNING-CLI.md @@ -45,10 +45,8 @@ before you consider going up a preset. You should try this parameter before touching the absolute values, since luma and chroma start from different defaults and the scale keeps that separation. -**`--lambda-ht` sets those thresholds outright.** The defaults, 5.3 for luma and 4.2 for chroma, -were tuned and deliberately biased toward keeping detail. A single value here flattens both planes -onto the same number, so prefer the scale unless you have a figure you want. -`--luma-lambda-ht` and `--chroma-lambda-ht` pin one plane without touching the other, and `--lambda-ht-scale` still +**`--lambda-ht` sets those thresholds outright.** The defaults are 5.2 for luma and 3.4 for chroma. +`--luma-lambda-ht` and `--chroma-lambda-ht` pin one plane without touching the other, and `--lambda-ht-scale` still applies on top of whatever is pinned. **`--sigma-scale` is the other one**, and it does something different. The lambda dials decide how @@ -61,6 +59,11 @@ are happy with the level and just want to adjust how much noise is removed vs de positions, so it dominates the work. Dropping it from 9 to 6 roughly halves the candidates, which is exactly what `--preset veryfast` does. +**`--field-lambda` pulls each block's motion vector toward its neighbours.** `1.0` (the library +default) is the shipped calibration and smooths the tracked field. Raise it further on noisy or +flat content where vectors wander, at the cost of following small objects less closely. `0` turns +the smoothing off and leaves the tracked field as it is. + ### What not to touch in NL4D - **`--sigma`** pins the noise level to a fixed value and turns the per-scene measurement off @@ -70,10 +73,12 @@ is exactly what `--preset veryfast` does. admitted once they are scored, so it is not a quality dial. - **`--no-confidence-variance`** stops a poorly matched patch from being trusted less than a well-matched one. It exists to isolate that mechanism in testing and calibration, not to improve output. -- **`--mismatch-scale`** sets how much less a poorly matched patch is trusted, rather than whether it is. +- **`--mismatch-scale`** sets how much less a poorly matched patch is trusted, rather than whether it + is, judged by the patch's own match residual rather than the motion block's score. The variance it controls grows with the square of the value, so `2` distrusts a bad match four times - as much. The effect saturates somewhere between `3` and `13` depending on how noisy the source is, - and `0` is the same thing as `--no-confidence-variance`. + as much. The effect saturates. It saturates sooner the worse the patch matched, because the + variance the mechanism derives is capped at 64 times the channel's own variance, and `0` is the + same thing as `--no-confidence-variance`. - **`--thsad-scale`, `--mc-blksize`, `--mc-overlap`, `--mc-search`, `--mc-pyramid-levels`** tune the motion machinery's internals, changing any of these will likely invalidate all other defaults. diff --git a/docs/TUNING-VS.md b/docs/TUNING-VS.md index 5edff69..68636cb 100644 --- a/docs/TUNING-VS.md +++ b/docs/TUNING-VS.md @@ -73,11 +73,10 @@ values, since luma and chroma start from different defaults and the scale keeps clean = avd.Nl4d(clip, lambda_ht_scale=1.1) ``` -**`lambda_ht` sets those thresholds outright.** The defaults, 5.3 for luma and 4.2 for chroma, -were tuned and deliberately biased toward keeping detail. A single value here flattens both planes -onto the same number, so prefer the scale unless you have a figure you want. `luma_lambda_ht` and -`chroma_lambda_ht`, both reachable by name, pin one plane without touching the other, and -`lambda_ht_scale` still applies on top of whatever is pinned. +**`lambda_ht` sets those thresholds outright.** The defaults are 5.2 for luma and 3.4 for chroma. +A single value here flattens both planes onto the same number, so prefer the scale unless you +have a figure you want. `luma_lambda_ht` and `chroma_lambda_ht`, both reachable by name, pin one +plane without touching the other, and `lambda_ht_scale` still applies on top of whatever is pinned. ```python clean = avd.Nl4d(clip, luma_lambda_ht=5.0, lambda_ht_scale=1.05) @@ -100,9 +99,11 @@ is exactly what `preset="veryfast"` does. entirely. `sigma_scale` keeps the measurement and nudges it, which is almost always what you actually want. - **`luma_mismatch_scale` and `chroma_mismatch_scale`** set how much less a poorly matched patch is - trusted, rather than whether it is. The variance they control grows with the square of the value, - so `2` distrusts a bad match four times as much. The effect saturates somewhere between `3` and - `13` depending on how noisy the source is, and `0` turns the mechanism off. + trusted, rather than whether it is, judged by the patch's own match residual rather than the + motion block's score. The variance they control grows with the square of the value, + so `2` distrusts a bad match four times as much. The effect saturates. It saturates sooner the + worse the patch matched, because the variance the mechanism derives is capped at 64 times the + channel's own variance, and `0` turns the mechanism off. - **`temporal_radius`** is what `preset` mostly exists to resolve. Setting it by hand is fine, but it is the same lever the preset ladder pulls, so reach for the ladder first. diff --git a/docs/TUTORIAL-CLI.md b/docs/TUTORIAL-CLI.md index 0020c33..1e5e810 100644 --- a/docs/TUTORIAL-CLI.md +++ b/docs/TUTORIAL-CLI.md @@ -156,7 +156,7 @@ What `--preset` fills in: | Flag | What it does | Default | |-------------------------|----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|----------------------| | `--temporal-radius ` | How many neighbouring frames to search on each side, between 1 and 8. More frames means more patches to group. | from `--preset` | -| `--lambda-ht ` | How aggressively small transform coefficients are zeroed out. Higher removes more noise and more fine detail with it. `--luma-lambda-ht` and `--chroma-lambda-ht` override one plane. | 5.3 luma, 4.2 chroma | +| `--lambda-ht ` | How aggressively small transform coefficients are zeroed out. Higher removes more noise and more fine detail with it. `--luma-lambda-ht` and `--chroma-lambda-ht` override one plane. | 5.2 luma, 3.4 chroma | | `--lambda-ht-scale ` | Multiplies the `--lambda-ht` in effect for each plane. The main quality dial, since luma and chroma start from different defaults and this moves both together. | `1.0` | | `--spatial-radius ` | Half-width of the candidate search inside the centre frame, between 1 and 16. Most of the search work goes here, since the window covers `(2N+1)^2` positions. | from `--preset` | | `--sigma-scale ` | Nudges the measured noise level, the same dial NLMeans spells `--hq-sigma-scale`. | `1.0` |