1use std::collections::VecDeque;
2
3use cubecl::Runtime;
4use cubecl::prelude::ComputeClient;
5
6use crate::accelerate::Accelerator;
7use crate::device::Device;
8use crate::nl4d::{Nl4dDenoiser, Nl4dParams};
9#[cfg(test)]
10use crate::nlmeans::MotionEstimation;
11use crate::nlmeans::{
12 ChannelMode,
13 HqParams,
14 MotionCompensationMode,
15 MotionSearch,
16 NlmDenoiser,
17 NlmParams,
18 Pending,
19 PrefilterMode,
20 hq_default_strength,
21 validate_dimensions,
22};
23use crate::sniff::sniff_best_accelerator;
24
25#[derive(Debug, Clone, bon::Builder)]
33pub struct DenoiserOptions {
34 #[builder(default = ChannelMode::Yuv)]
36 pub channel_mode: ChannelMode,
37 #[builder(default = DenoisingMode::Spacial)]
40 pub mode: DenoisingMode,
41 #[builder(default)]
44 pub algorithm: Algorithm,
45}
46
47#[derive(Debug, Copy, Clone, PartialEq)]
52pub enum Algorithm {
53 Nlmeans(NlmeansOptions),
56 NlmeansHq(NlmeansHqOptions),
62 Nl4d(Nl4dOptions),
69}
70
71impl Default for Algorithm {
72 fn default() -> Self {
73 Self::Nlmeans(NlmeansOptions::default())
74 }
75}
76
77#[derive(Debug, Copy, Clone, Default, PartialEq)]
79pub struct NlmeansOptions {
80 pub prefilter: PrefilterMode,
85 pub motion_compensation: MotionCompensationMode,
94 pub tuning: NlmTuning,
97}
98
99#[derive(Debug, Copy, Clone, Default, PartialEq)]
101pub struct NlmeansHqOptions {
102 pub nlm: NlmeansOptions,
104 pub hq: HqParams,
106}
107
108#[derive(Debug, Copy, Clone, PartialEq)]
126pub struct Nl4dOptions {
127 pub motion: MotionSearch,
129 pub sigma: Option<f32>,
135 pub sigma_scale: f32,
141 pub thsad_scale: f32,
147 pub refine: u32,
150 pub spatial_radius: u32,
153 pub lambda_ht: Option<f32>,
159 pub lambda_ht_scale: f32,
166 pub c_min: f32,
171 pub mismatch_scale: f32,
178 pub confidence_variance: bool,
182}
183
184impl Default for Nl4dOptions {
185 fn default() -> Self {
186 let defaults = Nl4dParams::default();
187 let hq = HqParams::default();
188 Self {
189 motion: MotionSearch::default(),
190 sigma: hq.sigma_override,
191 sigma_scale: hq.sigma_scale,
192 thsad_scale: hq.thsad_scale,
193 refine: defaults.refine,
194 spatial_radius: defaults.spatial_radius,
195 lambda_ht: None,
199 lambda_ht_scale: 1.0,
200 c_min: defaults.c_min,
201 mismatch_scale: defaults.mismatch_scale,
202 confidence_variance: defaults.confidence_variance,
203 }
204 }
205}
206
207impl Nl4dOptions {
208 fn to_hq_params(self) -> HqParams {
215 HqParams {
216 sigma_override: self.sigma,
217 sigma_scale: self.sigma_scale,
218 thsad_scale: self.thsad_scale,
219 temporal_confidence: true,
220 ..HqParams::default()
221 }
222 }
223}
224
225pub fn nl4d_default_lambda_ht(channels: ChannelMode) -> f32 {
244 match channels {
245 ChannelMode::Luma | ChannelMode::Yuv => 5.3,
246 ChannelMode::Chroma => 4.2,
247 }
248}
249
250fn resolve_lambda_ht(opts: &Nl4dOptions, channels: ChannelMode) -> Result<f32, String> {
262 if !(opts.lambda_ht_scale.is_finite() && (0.1..=10.0).contains(&opts.lambda_ht_scale)) {
263 return Err(format!(
264 "lambda_ht_scale must be finite and in [0.1, 10.0], got {}",
265 opts.lambda_ht_scale
266 ));
267 }
268
269 let lambda_ht = opts.lambda_ht.unwrap_or_else(|| nl4d_default_lambda_ht(channels));
270
271 Ok(lambda_ht * opts.lambda_ht_scale)
272}
273
274#[derive(Debug, Copy, Clone, Eq, PartialEq)]
276pub enum DenoisingMode {
277 Spacial,
279 Temporal { radius: u32 },
281}
282
283#[derive(Debug, Copy, Clone, Default, PartialEq)]
288pub struct NlmTuning {
289 pub search_radius: Option<u32>,
290 pub patch_radius: Option<u32>,
291 pub strength: Option<f32>,
292 pub self_weight: Option<f32>,
293}
294
295impl DenoiserOptions {
296 #[doc(hidden)]
310 pub fn to_nlm_params(&self) -> NlmParams {
311 let temporal_radius = match self.mode {
312 DenoisingMode::Spacial => 0,
313 DenoisingMode::Temporal { radius } => radius,
314 };
315
316 match self.algorithm {
317 Algorithm::Nlmeans(opts) => self.nlm_params_for(opts, None, temporal_radius),
318 Algorithm::NlmeansHq(opts) => self.nlm_params_for(opts.nlm, Some(opts.hq), temporal_radius),
319 Algorithm::Nl4d(opts) => NlmParams {
323 channels: self.channel_mode,
324 motion_compensation: opts.motion.into(),
325 temporal_radius,
326 hq: Some(opts.to_hq_params()),
327 ..NlmParams::default()
328 },
329 }
330 }
331
332 fn nlm_params_for(&self, opts: NlmeansOptions, hq: Option<HqParams>, temporal_radius: u32) -> NlmParams {
335 let strength = opts.tuning.strength.unwrap_or(match hq {
351 Some(hq) if hq.auto_strength => hq_default_strength(self.channel_mode, temporal_radius),
352 _ => NlmParams::default().strength,
353 });
354
355 let defaults = NlmParams::default();
356 NlmParams {
357 channels: self.channel_mode,
358 prefilter: opts.prefilter,
359 motion_compensation: opts.motion_compensation,
360 temporal_radius,
361 hq,
362 strength,
363 search_radius: opts.tuning.search_radius.unwrap_or(defaults.search_radius),
364 patch_radius: opts.tuning.patch_radius.unwrap_or(defaults.patch_radius),
365 self_weight: opts.tuning.self_weight.unwrap_or(defaults.self_weight),
366 }
367 }
368}
369
370#[derive(Debug, thiserror::Error)]
372pub enum DenoiserError {
373 #[error("denoiser queue is full, collect the pending frame before pushing more")]
379 QueueFull,
380 #[error("no accelerator from the priority list is available")]
382 NoAcceleratorAvailable,
383 #[error(transparent)]
386 Other(#[from] anyhow::Error),
387}
388
389enum Engine<R: Runtime> {
395 Nlm(Box<NlmDenoiser<R>>),
396 Nl4d(Box<Nl4dDenoiser<R>>),
397}
398
399impl<R: Runtime> Engine<R> {
400 fn push_frame(&mut self, frame: &[f32]) {
401 match self {
402 Self::Nlm(d) => d.push_frame(frame),
403 Self::Nl4d(d) => d.push_frame(frame),
404 }
405 }
406
407 fn denoise_submit(&mut self) -> Result<Option<Pending<R>>, anyhow::Error> {
408 match self {
409 Self::Nlm(d) => d.denoise_submit(),
410 Self::Nl4d(d) => d.denoise_submit().map_err(anyhow::Error::from),
415 }
416 }
417
418 fn flush(&mut self, sink: impl FnMut(&[f32])) -> Result<(), anyhow::Error> {
419 match self {
420 Self::Nlm(d) => d.flush(sink),
421 Self::Nl4d(d) => d.flush(sink).map_err(anyhow::Error::from),
422 }
423 }
424}
425
426fn build_engine<R: Runtime>(
436 client: &ComputeClient<R>,
437 algorithm: &Algorithm,
438 params: NlmParams,
439 width: u32,
440 height: u32,
441) -> Result<Engine<R>, DenoiserError> {
442 match algorithm {
443 Algorithm::Nl4d(opts) => {
444 if params.temporal_radius == 0 {
447 return Err(DenoiserError::Other(anyhow::anyhow!(
448 "nl4d needs a temporal window, set DenoiserOptions::mode to \
449 DenoisingMode::Temporal"
450 )));
451 }
452
453 let lambda_ht = resolve_lambda_ht(opts, params.channels)
454 .map_err(|e| DenoiserError::Other(anyhow::anyhow!(e)))?;
455 let nl4d_params = Nl4dParams {
456 temporal_radius: params.temporal_radius,
457 nlm: params,
458 refine: opts.refine,
459 spatial_radius: opts.spatial_radius,
460 lambda_ht,
461 c_min: opts.c_min,
462 mismatch_scale: opts.mismatch_scale,
463 confidence_variance: opts.confidence_variance,
464 };
465 let denoiser = Nl4dDenoiser::new(client, nl4d_params, width, height)
466 .map_err(|e| DenoiserError::Other(anyhow::anyhow!(e)))?;
467 Ok(Engine::Nl4d(Box::new(denoiser)))
468 },
469 Algorithm::Nlmeans(_) | Algorithm::NlmeansHq(_) => Ok(Engine::Nlm(Box::new(NlmDenoiser::new(
470 client, params, width, height,
471 )))),
472 }
473}
474
475enum Backend {
476 #[cfg(feature = "cuda")]
477 Cuda(Engine<cubecl::cuda::CudaRuntime>),
478 #[cfg(feature = "rocm")]
479 Rocm(Engine<cubecl::hip::HipRuntime>),
480 #[cfg(any(feature = "vulkan", feature = "metal"))]
481 Wgpu(Engine<cubecl::wgpu::WgpuRuntime>),
482}
483
484enum BackendPending {
485 #[cfg(feature = "cuda")]
486 Cuda(Pending<cubecl::cuda::CudaRuntime>),
487 #[cfg(feature = "rocm")]
488 Rocm(Pending<cubecl::hip::HipRuntime>),
489 #[cfg(any(feature = "vulkan", feature = "metal"))]
490 Wgpu(Pending<cubecl::wgpu::WgpuRuntime>),
491}
492
493impl BackendPending {
494 fn wait(self) -> Result<Vec<f32>, anyhow::Error> {
495 match self {
496 #[cfg(feature = "cuda")]
497 Self::Cuda(p) => p.wait(),
498 #[cfg(feature = "rocm")]
499 Self::Rocm(p) => p.wait(),
500 #[cfg(any(feature = "vulkan", feature = "metal"))]
501 Self::Wgpu(p) => p.wait(),
502 }
503 }
504}
505
506pub const MAX_PENDING: usize = 2;
513
514pub struct Denoiser {
564 backend: Backend,
565 pending: VecDeque<BackendPending>,
566 accelerator: Accelerator,
567 width: u32,
568 height: u32,
569 channels: u32,
570 temporal_radius: u32,
571 frames_pushed: u32,
572}
573
574impl Denoiser {
575 pub fn create(
602 accelerators: &[Accelerator],
603 device: &Device,
604 width: u32,
605 height: u32,
606 options: DenoiserOptions,
607 ) -> Result<Self, DenoiserError> {
608 let accelerator =
609 sniff_best_accelerator(accelerators, device).ok_or(DenoiserError::NoAcceleratorAvailable)?;
610
611 let params = options.to_nlm_params();
612 params.validate()?;
613 validate_dimensions(width, height)?;
614
615 let channels = params.channels.count();
616 let temporal_radius = params.temporal_radius;
617 let backend = build_backend(accelerator, device, &options.algorithm, params, width, height)?;
618
619 Ok(Self {
620 backend,
621 pending: VecDeque::with_capacity(MAX_PENDING),
622 accelerator,
623 width,
624 height,
625 channels,
626 temporal_radius,
627 frames_pushed: 0,
628 })
629 }
630
631 pub fn selected_accelerator(&self) -> Accelerator {
633 self.accelerator
634 }
635
636 pub fn width(&self) -> u32 {
638 self.width
639 }
640
641 pub fn height(&self) -> u32 {
643 self.height
644 }
645
646 pub fn push_frame(&mut self, frame: &[f32]) -> Result<(), DenoiserError> {
660 let window_full = self.frames_pushed > self.temporal_radius;
664 if window_full && self.pending.len() >= MAX_PENDING {
665 return Err(DenoiserError::QueueFull);
666 }
667
668 match &mut self.backend {
669 #[cfg(feature = "cuda")]
670 Backend::Cuda(d) => {
671 d.push_frame(frame);
672 if let Some(p) = d.denoise_submit()? {
673 self.pending.push_back(BackendPending::Cuda(p));
674 }
675 },
676 #[cfg(feature = "rocm")]
677 Backend::Rocm(d) => {
678 d.push_frame(frame);
679 if let Some(p) = d.denoise_submit()? {
680 self.pending.push_back(BackendPending::Rocm(p));
681 }
682 },
683 #[cfg(any(feature = "vulkan", feature = "metal"))]
684 Backend::Wgpu(d) => {
685 d.push_frame(frame);
686 if let Some(p) = d.denoise_submit()? {
687 self.pending.push_back(BackendPending::Wgpu(p));
688 }
689 },
690 }
691
692 self.frames_pushed = self.frames_pushed.saturating_add(1);
693 Ok(())
694 }
695
696 pub fn recv_frame(&mut self) -> Result<Option<Vec<f32>>, DenoiserError> {
702 let Some(pending) = self.pending.pop_front() else {
703 return Ok(None);
704 };
705 Ok(Some(pending.wait()?))
706 }
707
708 pub fn try_recv_frame(&mut self) -> Result<Option<Vec<f32>>, DenoiserError> {
714 self.recv_frame()
715 }
716
717 pub fn flush(&mut self, mut sink: impl FnMut(Vec<f32>)) -> Result<(), DenoiserError> {
730 while let Some(frame) = self.recv_frame()? {
733 sink(frame);
734 }
735
736 let pixels = (self.width * self.height) as usize;
737 let channels = self.channels as usize;
738 let scratch_cap = pixels * channels;
739
740 match &mut self.backend {
741 #[cfg(feature = "cuda")]
742 Backend::Cuda(d) => d.flush(|slice| {
743 let mut v = Vec::with_capacity(scratch_cap);
744 v.extend_from_slice(slice);
745 sink(v);
746 })?,
747 #[cfg(feature = "rocm")]
748 Backend::Rocm(d) => d.flush(|slice| {
749 let mut v = Vec::with_capacity(scratch_cap);
750 v.extend_from_slice(slice);
751 sink(v);
752 })?,
753 #[cfg(any(feature = "vulkan", feature = "metal"))]
754 Backend::Wgpu(d) => d.flush(|slice| {
755 let mut v = Vec::with_capacity(scratch_cap);
756 v.extend_from_slice(slice);
757 sink(v);
758 })?,
759 }
760
761 self.frames_pushed = 0;
765
766 Ok(())
767 }
768}
769
770fn build_backend(
771 accel: Accelerator,
772 device: &Device,
773 algorithm: &Algorithm,
774 params: NlmParams,
775 width: u32,
776 height: u32,
777) -> Result<Backend, DenoiserError> {
778 match accel {
779 #[cfg(feature = "cuda")]
780 Accelerator::Cuda => {
781 let dev = device.to_cuda()?;
782 let client = <cubecl::cuda::CudaRuntime as Runtime>::client(&dev);
783 Ok(Backend::Cuda(build_engine(
784 &client, algorithm, params, width, height,
785 )?))
786 },
787 #[cfg(feature = "rocm")]
788 Accelerator::Rocm => {
789 let dev = device.to_amd()?;
790 let client = <cubecl::hip::HipRuntime as Runtime>::client(&dev);
791 Ok(Backend::Rocm(build_engine(
792 &client, algorithm, params, width, height,
793 )?))
794 },
795 #[cfg(feature = "vulkan")]
796 Accelerator::Vulkan => {
797 let dev = device.to_wgpu()?;
798 let client = <cubecl::wgpu::WgpuRuntime as Runtime>::client(&dev);
799 Ok(Backend::Wgpu(build_engine(
800 &client, algorithm, params, width, height,
801 )?))
802 },
803 #[cfg(feature = "metal")]
804 Accelerator::Metal => {
805 let dev = device.to_wgpu()?;
806 let client = <cubecl::wgpu::WgpuRuntime as Runtime>::client(&dev);
807 Ok(Backend::Wgpu(build_engine(
808 &client, algorithm, params, width, height,
809 )?))
810 },
811 #[cfg(docsrs)]
815 #[allow(unreachable_patterns)]
816 _ => unreachable!(),
817 }
818}
819
820#[cfg(test)]
821mod options_tests {
822 use super::*;
823
824 fn hq(hq: HqParams) -> Algorithm {
827 Algorithm::NlmeansHq(NlmeansHqOptions {
828 hq,
829 ..NlmeansHqOptions::default()
830 })
831 }
832
833 fn fast_tuned(tuning: NlmTuning) -> Algorithm {
835 Algorithm::Nlmeans(NlmeansOptions {
836 tuning,
837 ..NlmeansOptions::default()
838 })
839 }
840
841 #[test]
842 fn nl4d_default_lambda_ht_differs_between_luma_and_chroma() {
843 let luma = nl4d_default_lambda_ht(ChannelMode::Luma);
844 let chroma = nl4d_default_lambda_ht(ChannelMode::Chroma);
845
846 assert!((luma - 5.3).abs() < f32::EPSILON);
847 assert!((chroma - 4.2).abs() < f32::EPSILON);
848 assert!(
849 (chroma - luma).abs() > f32::EPSILON,
850 "the two planes should not resolve to the same default"
851 );
852 }
853
854 #[test]
855 fn nl4d_default_lambda_ht_yuv_reads_the_luma_value() {
856 let yuv = nl4d_default_lambda_ht(ChannelMode::Yuv);
857 let luma = nl4d_default_lambda_ht(ChannelMode::Luma);
858
859 assert!((yuv - luma).abs() < f32::EPSILON);
860 }
861
862 #[test]
863 fn resolve_lambda_ht_unset_uses_the_per_plane_default() {
864 let opts = Nl4dOptions::default();
865
866 let luma = resolve_lambda_ht(&opts, ChannelMode::Luma).expect("the default scale is in range");
867 let chroma = resolve_lambda_ht(&opts, ChannelMode::Chroma).expect("the default scale is in range");
868
869 assert!((luma - 5.3).abs() < f32::EPSILON, "got {luma}");
870 assert!((chroma - 4.2).abs() < f32::EPSILON, "got {chroma}");
871 }
872
873 #[test]
874 fn resolve_lambda_ht_explicit_value_overrides_every_plane() {
875 let opts = Nl4dOptions {
876 lambda_ht: Some(4.4),
877 ..Nl4dOptions::default()
878 };
879
880 for channels in [ChannelMode::Luma, ChannelMode::Chroma, ChannelMode::Yuv] {
881 let got = resolve_lambda_ht(&opts, channels).expect("the default scale is in range");
882 assert!(
883 (got - 4.4).abs() < f32::EPSILON,
884 "channels {channels:?} got {got}"
885 );
886 }
887 }
888
889 #[test]
890 fn resolve_lambda_ht_default_scale_leaves_the_value_alone() {
891 let opts = Nl4dOptions::default();
892
893 for channels in [ChannelMode::Luma, ChannelMode::Chroma, ChannelMode::Yuv] {
894 let got = resolve_lambda_ht(&opts, channels).expect("the default scale is in range");
895 let want = nl4d_default_lambda_ht(channels);
896 assert!(
897 (got - want).abs() < f32::EPSILON,
898 "channels {channels:?} got {got}"
899 );
900 }
901 }
902
903 #[test]
904 fn resolve_lambda_ht_scale_multiplies_the_per_plane_default() {
905 let opts = Nl4dOptions {
906 lambda_ht_scale: 1.1,
907 ..Nl4dOptions::default()
908 };
909
910 for channels in [ChannelMode::Luma, ChannelMode::Chroma, ChannelMode::Yuv] {
911 let got = resolve_lambda_ht(&opts, channels).expect("1.1 is in range");
912 let want = nl4d_default_lambda_ht(channels) * 1.1;
913 assert!(
914 (got - want).abs() < 1e-5,
915 "channels {channels:?} got {got}, want {want}"
916 );
917 }
918 }
919
920 #[test]
923 fn resolve_lambda_ht_scale_multiplies_an_explicit_value() {
924 let opts = Nl4dOptions {
925 lambda_ht: Some(5.0),
926 lambda_ht_scale: 0.9,
927 ..Nl4dOptions::default()
928 };
929
930 let got = resolve_lambda_ht(&opts, ChannelMode::Luma).expect("0.9 is in range");
931 assert!((got - 4.5).abs() < 1e-5, "got {got}");
932 }
933
934 #[test]
935 fn resolve_lambda_ht_rejects_an_out_of_range_scale() {
936 for bad in [0.0, -1.0, 0.05, 10.5, f32::NAN, f32::INFINITY] {
937 let opts = Nl4dOptions {
938 lambda_ht_scale: bad,
939 ..Nl4dOptions::default()
940 };
941 let err = resolve_lambda_ht(&opts, ChannelMode::Luma).unwrap_err();
942 assert!(
943 err.contains("lambda_ht_scale"),
944 "lambda_ht_scale={bad} should be rejected, got {err}"
945 );
946 }
947 }
948
949 #[test]
950 fn the_default_algorithm_is_the_fast_nlmeans_path() {
951 let opts = DenoiserOptions::builder().build();
952 assert_eq!(opts.algorithm, Algorithm::Nlmeans(NlmeansOptions::default()));
953 }
954
955 #[test]
956 fn spatial_mode_maps_to_zero_temporal_radius() {
957 let opts = DenoiserOptions::builder()
958 .channel_mode(ChannelMode::Yuv)
959 .mode(DenoisingMode::Spacial)
960 .build();
961 let params = opts.to_nlm_params();
962
963 assert_eq!(params.temporal_radius, 0);
964 assert_eq!(params.channels, ChannelMode::Yuv);
965 }
966
967 #[test]
968 fn temporal_mode_propagates_radius() {
969 let opts = DenoiserOptions::builder()
970 .mode(DenoisingMode::Temporal { radius: 3 })
971 .build();
972 let params = opts.to_nlm_params();
973
974 assert_eq!(params.temporal_radius, 3);
975 }
976
977 #[test]
978 fn prefilter_passthrough() {
979 let opts = DenoiserOptions::builder()
980 .algorithm(Algorithm::Nlmeans(NlmeansOptions {
981 prefilter: PrefilterMode::Bilateral {
982 sigma_s: 3.0,
983 sigma_r: 0.02,
984 },
985 ..NlmeansOptions::default()
986 }))
987 .build();
988 let params = opts.to_nlm_params();
989
990 assert!(matches!(params.prefilter, PrefilterMode::Bilateral { .. }));
991 }
992
993 #[test]
994 fn hq_unset_prefilter_defaults_to_none() {
995 let opts = DenoiserOptions::builder()
996 .algorithm(hq(HqParams::default()))
997 .build();
998 let params = opts.to_nlm_params();
999
1000 assert!(matches!(params.prefilter, PrefilterMode::None));
1001 }
1002
1003 #[test]
1004 fn fast_unset_prefilter_defaults_to_none() {
1005 let opts = DenoiserOptions::builder()
1006 .algorithm(Algorithm::Nlmeans(NlmeansOptions::default()))
1007 .build();
1008 let params = opts.to_nlm_params();
1009
1010 assert!(matches!(params.prefilter, PrefilterMode::None));
1011 }
1012
1013 #[test]
1014 fn hq_unset_strength_defaults_to_hq_default_strength() {
1015 let opts = DenoiserOptions::builder()
1017 .algorithm(hq(HqParams::default()))
1018 .build();
1019 let params = opts.to_nlm_params();
1020
1021 let expected = hq_default_strength(ChannelMode::Yuv, 0);
1022 assert!((params.strength - expected).abs() < f32::EPSILON);
1023 }
1024
1025 #[test]
1026 fn hq_no_auto_strength_falls_back_to_the_legacy_absolute_default() {
1027 let opts = DenoiserOptions::builder()
1034 .algorithm(hq(HqParams {
1035 auto_strength: false,
1036 ..HqParams::default()
1037 }))
1038 .build();
1039 let params = opts.to_nlm_params();
1040
1041 let expected = NlmParams::default().strength;
1042 assert!(
1043 (params.strength - expected).abs() < f32::EPSILON,
1044 "expected the legacy absolute default {expected}, got {}, which looks like the \
1045 auto-strength multiplier table leaking through",
1046 params.strength
1047 );
1048 }
1049
1050 #[test]
1051 fn hq_luma_r4_uses_measured_table_value() {
1052 let opts = DenoiserOptions::builder()
1053 .channel_mode(ChannelMode::Luma)
1054 .mode(DenoisingMode::Temporal { radius: 4 })
1055 .algorithm(hq(HqParams::default()))
1056 .build();
1057 let params = opts.to_nlm_params();
1058
1059 assert!((params.strength - 0.35).abs() < f32::EPSILON);
1060 }
1061
1062 #[test]
1063 fn hq_chroma_r4_uses_measured_table_value() {
1064 let opts = DenoiserOptions::builder()
1065 .channel_mode(ChannelMode::Chroma)
1066 .mode(DenoisingMode::Temporal { radius: 4 })
1067 .algorithm(hq(HqParams::default()))
1068 .build();
1069 let params = opts.to_nlm_params();
1070
1071 assert!((params.strength - 0.70).abs() < f32::EPSILON);
1072 }
1073
1074 #[test]
1075 fn hq_yuv_r8_uses_measured_table_value() {
1076 let opts = DenoiserOptions::builder()
1077 .channel_mode(ChannelMode::Yuv)
1078 .mode(DenoisingMode::Temporal { radius: 8 })
1079 .algorithm(hq(HqParams::default()))
1080 .build();
1081 let params = opts.to_nlm_params();
1082
1083 assert!((params.strength - 0.30).abs() < f32::EPSILON);
1084 }
1085
1086 #[test]
1087 fn hq_spacial_mode_uses_radius_zero_table_values() {
1088 for channels in [ChannelMode::Luma, ChannelMode::Chroma, ChannelMode::Yuv] {
1089 let opts = DenoiserOptions::builder()
1090 .channel_mode(channels)
1091 .mode(DenoisingMode::Spacial)
1092 .algorithm(hq(HqParams::default()))
1093 .build();
1094 let params = opts.to_nlm_params();
1095
1096 let expected = hq_default_strength(channels, 0);
1097 assert!(
1098 (params.strength - expected).abs() < f32::EPSILON,
1099 "for channels {channels:?} expected {expected}, got {}",
1100 params.strength
1101 );
1102 }
1103 }
1104
1105 #[test]
1106 fn hq_explicit_strength_wins_over_the_table_for_every_plane() {
1107 for channels in [ChannelMode::Luma, ChannelMode::Chroma, ChannelMode::Yuv] {
1108 let opts = DenoiserOptions::builder()
1109 .channel_mode(channels)
1110 .mode(DenoisingMode::Temporal { radius: 4 })
1111 .algorithm(Algorithm::NlmeansHq(NlmeansHqOptions {
1112 nlm: NlmeansOptions {
1113 tuning: NlmTuning {
1114 strength: Some(0.99),
1115 ..NlmTuning::default()
1116 },
1117 ..NlmeansOptions::default()
1118 },
1119 hq: HqParams::default(),
1120 }))
1121 .build();
1122 let params = opts.to_nlm_params();
1123
1124 assert!(
1125 (params.strength - 0.99).abs() < f32::EPSILON,
1126 "for channels {channels:?} the explicit strength was overridden by the table"
1127 );
1128 }
1129 }
1130
1131 #[test]
1132 fn fast_unset_strength_defaults_to_legacy_default() {
1133 let opts = DenoiserOptions::builder()
1134 .algorithm(Algorithm::Nlmeans(NlmeansOptions::default()))
1135 .build();
1136 let params = opts.to_nlm_params();
1137
1138 assert!((params.strength - 1.2).abs() < f32::EPSILON);
1139 }
1140
1141 #[test]
1142 fn nl4d_options_default_matches_nl4d_params_default() {
1143 let opts = Nl4dOptions::default();
1144 let params = crate::nl4d::Nl4dParams::default();
1145
1146 assert_eq!(opts.refine, params.refine);
1147 assert_eq!(opts.spatial_radius, params.spatial_radius);
1148 assert!((opts.c_min - params.c_min).abs() < f32::EPSILON);
1149 assert_eq!(opts.confidence_variance, params.confidence_variance);
1150 assert_eq!(opts.lambda_ht, None);
1157 assert!((params.lambda_ht - nl4d_default_lambda_ht(ChannelMode::Yuv)).abs() < f32::EPSILON);
1158 }
1159
1160 #[test]
1164 fn nl4d_builds_the_front_ends_hq_params_from_its_own_fields() {
1165 let opts = DenoiserOptions::builder()
1166 .mode(DenoisingMode::Temporal { radius: 2 })
1167 .algorithm(Algorithm::Nl4d(Nl4dOptions {
1168 sigma: Some(0.02),
1169 sigma_scale: 1.3,
1170 thsad_scale: 0.8,
1171 ..Nl4dOptions::default()
1172 }))
1173 .build();
1174 let params = opts.to_nlm_params();
1175
1176 let hq = params.hq.expect("nl4d always runs the hq front end");
1177 assert_eq!(hq.sigma_override, Some(0.02));
1178 assert!((hq.sigma_scale - 1.3).abs() < f32::EPSILON);
1179 assert!((hq.thsad_scale - 0.8).abs() < f32::EPSILON);
1180 assert!(
1181 hq.temporal_confidence,
1182 "the grouping kernel reads the confidence scores, so this cannot be off"
1183 );
1184 }
1185
1186 #[test]
1189 fn nl4d_reads_its_temporal_radius_from_the_denoising_mode() {
1190 for radius in [1u32, 4, 8] {
1191 let opts = DenoiserOptions::builder()
1192 .mode(DenoisingMode::Temporal { radius })
1193 .algorithm(Algorithm::Nl4d(Nl4dOptions::default()))
1194 .build();
1195
1196 assert_eq!(opts.to_nlm_params().temporal_radius, radius);
1197 }
1198 }
1199
1200 #[test]
1203 fn nl4d_never_builds_a_prefilter() {
1204 let opts = DenoiserOptions::builder()
1205 .mode(DenoisingMode::Temporal { radius: 2 })
1206 .algorithm(Algorithm::Nl4d(Nl4dOptions::default()))
1207 .build();
1208
1209 assert!(matches!(opts.to_nlm_params().prefilter, PrefilterMode::None));
1210 }
1211
1212 #[test]
1215 fn nl4d_leaves_the_nlm_weighting_knobs_at_their_defaults() {
1216 let defaults = NlmParams::default();
1217 let opts = DenoiserOptions::builder()
1218 .channel_mode(ChannelMode::Luma)
1219 .mode(DenoisingMode::Temporal { radius: 4 })
1220 .algorithm(Algorithm::Nl4d(Nl4dOptions::default()))
1221 .build();
1222 let params = opts.to_nlm_params();
1223
1224 assert!((params.strength - defaults.strength).abs() < f32::EPSILON);
1225 assert_eq!(params.search_radius, defaults.search_radius);
1226 assert_eq!(params.patch_radius, defaults.patch_radius);
1227 assert!((params.self_weight - defaults.self_weight).abs() < f32::EPSILON);
1228 }
1229
1230 #[test]
1233 fn nl4d_motion_search_becomes_an_active_mvtools_mode() {
1234 let opts = DenoiserOptions::builder()
1235 .mode(DenoisingMode::Temporal { radius: 2 })
1236 .algorithm(Algorithm::Nl4d(Nl4dOptions {
1237 motion: MotionSearch {
1238 blksize: 32,
1239 overlap: 16,
1240 search_radius: 6,
1241 pyramid_levels: 1,
1242 estimation: MotionEstimation::Direct,
1243 },
1244 ..Nl4dOptions::default()
1245 }))
1246 .build();
1247 let params = opts.to_nlm_params();
1248
1249 assert!(matches!(
1250 params.motion_compensation,
1251 MotionCompensationMode::Mvtools {
1252 blksize: 32,
1253 overlap: 16,
1254 search_radius: 6,
1255 pyramid_levels: 1,
1256 estimation: MotionEstimation::Direct,
1257 }
1258 ));
1259 }
1260
1261 #[test]
1262 fn nl4d_motion_search_defaults_match_the_front_ends_own_defaults() {
1263 let opts = DenoiserOptions::builder()
1264 .mode(DenoisingMode::Temporal { radius: 2 })
1265 .algorithm(Algorithm::Nl4d(Nl4dOptions::default()))
1266 .build();
1267 let params = opts.to_nlm_params();
1268
1269 assert_eq!(
1270 params.motion_compensation,
1271 crate::nl4d::Nl4dParams::default().nlm.motion_compensation
1272 );
1273 }
1274
1275 #[test]
1276 fn motion_compensation_passthrough() {
1277 let opts = DenoiserOptions::builder()
1278 .mode(DenoisingMode::Temporal { radius: 1 })
1279 .algorithm(Algorithm::Nlmeans(NlmeansOptions {
1280 motion_compensation: MotionCompensationMode::Mvtools {
1281 blksize: 16,
1282 overlap: 8,
1283 search_radius: 4,
1284 pyramid_levels: 2,
1285 estimation: MotionEstimation::Direct,
1286 },
1287 ..NlmeansOptions::default()
1288 }))
1289 .build();
1290 let params = opts.to_nlm_params();
1291
1292 assert!(matches!(
1293 params.motion_compensation,
1294 MotionCompensationMode::Mvtools {
1295 blksize: 16,
1296 overlap: 8,
1297 search_radius: 4,
1298 pyramid_levels: 2,
1299 ..
1300 }
1301 ));
1302 }
1303
1304 #[test]
1305 fn motion_compensation_defaults_to_none() {
1306 let opts = DenoiserOptions::builder().build();
1307 let params = opts.to_nlm_params();
1308 assert!(matches!(params.motion_compensation, MotionCompensationMode::None));
1309 }
1310
1311 #[test]
1312 fn nlm_tuning_overrides_individual_fields() {
1313 let defaults = NlmParams::default();
1314 let opts = DenoiserOptions::builder()
1315 .algorithm(fast_tuned(NlmTuning {
1316 search_radius: Some(7),
1317 patch_radius: None,
1318 strength: Some(2.5),
1319 self_weight: None,
1320 }))
1321 .build();
1322 let params = opts.to_nlm_params();
1323
1324 assert_eq!(params.search_radius, 7);
1325 assert_eq!(params.patch_radius, defaults.patch_radius);
1326 assert!((params.strength - 2.5).abs() < f32::EPSILON);
1327 assert!((params.self_weight - defaults.self_weight).abs() < f32::EPSILON);
1328 }
1329}
1330
1331#[cfg(all(test, feature = "vulkan"))]
1332mod tests {
1333 use super::*;
1334
1335 fn opts(mode: DenoisingMode) -> DenoiserOptions {
1336 DenoiserOptions::builder()
1337 .channel_mode(ChannelMode::Luma)
1338 .mode(mode)
1339 .build()
1340 }
1341
1342 fn frame(w: u32, h: u32) -> Vec<f32> {
1343 vec![0.5f32; (w * h) as usize]
1344 }
1345
1346 #[test]
1347 fn spatial_denoise_roundtrip() {
1348 let mut d = Denoiser::create(
1349 &[Accelerator::Vulkan],
1350 &Device::Default,
1351 16,
1352 16,
1353 opts(DenoisingMode::Spacial),
1354 )
1355 .expect("denoiser construction failed");
1356 assert_eq!(d.selected_accelerator(), Accelerator::Vulkan);
1357
1358 d.push_frame(&frame(16, 16)).expect("push failed");
1359 let out = d.recv_frame().expect("recv failed").expect("no frame");
1360 assert_eq!(out.len(), 16 * 16);
1361 }
1362
1363 #[test]
1364 fn nl4d_algorithm_round_trips_through_the_facade() {
1365 let opts = DenoiserOptions::builder()
1366 .channel_mode(ChannelMode::Luma)
1367 .mode(DenoisingMode::Temporal { radius: 2 })
1368 .algorithm(Algorithm::Nl4d(Nl4dOptions::default()))
1369 .build();
1370 let mut d = Denoiser::create(&[Accelerator::Vulkan], &Device::Default, 16, 16, opts)
1371 .expect("nl4d denoiser construction failed");
1372 assert_eq!(d.selected_accelerator(), Accelerator::Vulkan);
1373
1374 d.push_frame(&frame(16, 16)).expect("push failed");
1378 assert!(d.recv_frame().expect("recv failed").is_none());
1379
1380 let mut out = Vec::new();
1381 d.flush(|f| out.push(f)).expect("flush failed");
1382 assert_eq!(out.len(), 1, "expected exactly one output for one pushed frame");
1383 assert_eq!(out[0].len(), 16 * 16);
1384 }
1385
1386 #[test]
1390 fn nl4d_rejects_a_spatial_denoising_mode() {
1391 let opts = DenoiserOptions::builder()
1392 .channel_mode(ChannelMode::Luma)
1393 .mode(DenoisingMode::Spacial)
1394 .algorithm(Algorithm::Nl4d(Nl4dOptions::default()))
1395 .build();
1396 let result = Denoiser::create(&[Accelerator::Vulkan], &Device::Default, 16, 16, opts);
1397
1398 match result {
1399 Err(DenoiserError::Other(e)) => assert!(
1400 e.to_string().contains("temporal window"),
1401 "unexpected error message: {e}"
1402 ),
1403 Err(other) => panic!("expected DenoiserError::Other, got {other:?}"),
1404 Ok(_) => panic!("expected a rejection, got Ok"),
1405 }
1406 }
1407
1408 #[test]
1409 fn invalid_params_surface_as_error() {
1410 let bad = DenoiserOptions::builder()
1411 .algorithm(Algorithm::Nlmeans(NlmeansOptions {
1412 tuning: NlmTuning {
1413 strength: Some(0.0),
1414 ..NlmTuning::default()
1415 },
1416 ..NlmeansOptions::default()
1417 }))
1418 .build();
1419 let result = Denoiser::create(&[Accelerator::Vulkan], &Device::Default, 16, 16, bad);
1420
1421 match result {
1422 Err(DenoiserError::Other(_)) => {},
1423 Err(other) => panic!("expected DenoiserError::Other, got {other:?}"),
1424 Ok(_) => panic!("expected validation error, got Ok"),
1425 }
1426 }
1427
1428 #[test]
1429 fn tiny_frame_dimensions_surface_as_error() {
1430 let result = Denoiser::create(
1431 &[Accelerator::Vulkan],
1432 &Device::Default,
1433 2,
1434 2,
1435 opts(DenoisingMode::Spacial),
1436 );
1437
1438 match result {
1439 Err(DenoiserError::Other(e)) => {
1440 assert!(
1441 e.to_string().contains("supported minimum"),
1442 "unexpected error message: {e}"
1443 );
1444 },
1445 Err(other) => panic!("expected DenoiserError::Other, got {other:?}"),
1446 Ok(_) => panic!("expected dimension validation error, got Ok"),
1447 }
1448 }
1449
1450 #[test]
1451 fn push_after_pending_returns_queue_full() {
1452 let mut d = Denoiser::create(
1453 &[Accelerator::Vulkan],
1454 &Device::Default,
1455 16,
1456 16,
1457 opts(DenoisingMode::Spacial),
1458 )
1459 .unwrap();
1460
1461 d.push_frame(&frame(16, 16)).unwrap();
1466 d.push_frame(&frame(16, 16)).unwrap();
1467 let err = d.push_frame(&frame(16, 16)).expect_err("expected QueueFull");
1468 assert!(matches!(err, DenoiserError::QueueFull));
1469
1470 let out = d.recv_frame().unwrap().unwrap();
1471 assert_eq!(out.len(), 16 * 16);
1472
1473 d.push_frame(&frame(16, 16)).expect("push after drain failed");
1475 }
1476
1477 fn frame_filled(w: u32, h: u32, value: f32) -> Vec<f32> {
1478 vec![value; (w * h) as usize]
1479 }
1480
1481 fn push_n_with_drain(d: &mut Denoiser, n: usize, value: f32, out: &mut Vec<Vec<f32>>) {
1484 for _ in 0..n {
1485 loop {
1486 match d.push_frame(&frame_filled(16, 16, value)) {
1487 Ok(()) => break,
1488 Err(DenoiserError::QueueFull) => {
1489 let f = d
1490 .recv_frame()
1491 .expect("recv ok")
1492 .expect("queue full but recv yielded none");
1493 out.push(f);
1494 },
1495 Err(e) => panic!("unexpected push error: {e:?}"),
1496 }
1497 }
1498 }
1499 }
1500
1501 #[test]
1502 fn flush_leaves_denoiser_reusable_spatial() {
1503 let mut d = Denoiser::create(
1504 &[Accelerator::Vulkan],
1505 &Device::Default,
1506 16,
1507 16,
1508 opts(DenoisingMode::Spacial),
1509 )
1510 .unwrap();
1511
1512 let mut batch_a = Vec::new();
1513 push_n_with_drain(&mut d, 5, 0.25, &mut batch_a);
1514 d.flush(|f| batch_a.push(f)).expect("first flush failed");
1515 assert_eq!(batch_a.len(), 5);
1516
1517 assert!(d.recv_frame().unwrap().is_none());
1519
1520 let mut batch_b = Vec::new();
1521 push_n_with_drain(&mut d, 5, 0.75, &mut batch_b);
1522 d.flush(|f| batch_b.push(f)).expect("second flush failed");
1523 assert_eq!(batch_b.len(), 5);
1524
1525 for v in batch_b.iter().flatten() {
1526 assert!((v - 0.75).abs() < 0.1, "batch_b carried state from batch_a: {v}");
1527 }
1528 for v in batch_a.iter().flatten() {
1529 assert!((v - 0.25).abs() < 0.1, "batch_a value unexpectedly drifted: {v}");
1530 }
1531 }
1532
1533 #[test]
1534 fn flush_leaves_denoiser_reusable_temporal() {
1535 let mut d = Denoiser::create(
1536 &[Accelerator::Vulkan],
1537 &Device::Default,
1538 16,
1539 16,
1540 opts(DenoisingMode::Temporal { radius: 1 }),
1541 )
1542 .unwrap();
1543
1544 let mut batch_a = Vec::new();
1545 push_n_with_drain(&mut d, 5, 0.25, &mut batch_a);
1546 d.flush(|f| batch_a.push(f)).expect("first flush failed");
1547 assert_eq!(batch_a.len(), 5, "expected 5 frames from first batch");
1548
1549 assert!(d.recv_frame().unwrap().is_none());
1554 d.push_frame(&frame_filled(16, 16, 0.75)).unwrap();
1555 assert!(
1556 d.recv_frame().unwrap().is_none(),
1557 "first push of new temporal stream should not produce output yet"
1558 );
1559
1560 let mut batch_b = Vec::new();
1562 push_n_with_drain(&mut d, 4, 0.75, &mut batch_b);
1563 d.flush(|f| batch_b.push(f)).expect("second flush failed");
1564 assert_eq!(batch_b.len(), 5, "expected 5 frames from second batch");
1565
1566 for v in batch_b.iter().flatten() {
1567 assert!((v - 0.75).abs() < 0.1, "batch_b carried state from batch_a: {v}");
1568 }
1569 }
1570
1571 #[test]
1572 fn flush_emits_exactly_n_outputs_for_small_n() {
1573 for n in 1..=5usize {
1578 let mut d = Denoiser::create(
1579 &[Accelerator::Vulkan],
1580 &Device::Default,
1581 16,
1582 16,
1583 opts(DenoisingMode::Temporal { radius: 2 }),
1584 )
1585 .unwrap();
1586
1587 let mut out = Vec::new();
1588 push_n_with_drain(&mut d, n, 0.5, &mut out);
1589 d.flush(|f| out.push(f)).expect("flush failed");
1590 assert_eq!(
1591 out.len(),
1592 n,
1593 "expected {n} outputs for {n} pushes, got {}",
1594 out.len()
1595 );
1596 }
1597 }
1598}