1use std::collections::VecDeque;
2
3use crate::accelerate::Accelerator;
4use crate::{
5 Algorithm,
6 ChannelMode,
7 Denoiser,
8 DenoiserError,
9 DenoiserOptions,
10 DenoisingMode,
11 Depth,
12 Device,
13 Nl4dOptions,
14 NlmTuning,
15 NlmeansHqOptions,
16 NlmeansOptions,
17 WindowSpan,
18};
19
20#[derive(Debug, Clone, Copy, PartialEq, Eq)]
21pub enum Subsampling {
22 Yuv420,
23 Yuv422,
24 Yuv444,
25}
26
27impl Subsampling {
28 pub fn chroma_dims(self, w: u32, h: u32) -> (u32, u32) {
29 match self {
30 Subsampling::Yuv420 => (w / 2, h / 2),
31 Subsampling::Yuv422 => (w / 2, h),
32 Subsampling::Yuv444 => (w, h),
33 }
34 }
35}
36
37#[derive(Debug, Clone, Copy)]
38pub struct FrameLayout {
39 pub width: u32,
40 pub height: u32,
41 pub subsampling: Subsampling,
42 pub depth: Depth,
43}
44
45impl FrameLayout {
46 pub fn luma_pixels(&self) -> usize {
47 (self.width as usize) * (self.height as usize)
48 }
49
50 pub fn chroma_dims(&self) -> (u32, u32) {
51 self.subsampling.chroma_dims(self.width, self.height)
52 }
53
54 pub fn chroma_pixels(&self) -> usize {
55 let (w, h) = self.chroma_dims();
56 (w as usize) * (h as usize)
57 }
58
59 pub fn luma_bytes(&self) -> usize {
61 self.luma_pixels() * self.depth.bytes_per_sample()
62 }
63
64 pub fn chroma_bytes(&self) -> usize {
66 self.chroma_pixels() * self.depth.bytes_per_sample()
67 }
68
69 pub fn black_luma_plane(&self) -> Vec<u8> {
71 fill_plane(self.luma_pixels(), 0, self.depth)
72 }
73
74 pub fn neutral_chroma_plane(&self) -> Vec<u8> {
76 fill_plane(self.chroma_pixels(), self.depth.neutral_chroma(), self.depth)
77 }
78}
79
80pub fn fill_plane(samples: usize, value: u16, depth: Depth) -> Vec<u8> {
82 match depth.bytes_per_sample() {
83 1 => vec![value as u8; samples],
84 _ => {
85 let word = value.to_le_bytes();
86 let mut out = Vec::with_capacity(samples * 2);
87 for _ in 0..samples {
88 out.extend_from_slice(&word);
89 }
90 out
91 },
92 }
93}
94
95#[derive(Debug, Clone)]
101pub struct Planes {
102 pub y: Vec<u8>,
103 pub u: Vec<u8>,
104 pub v: Vec<u8>,
105}
106
107#[derive(Debug, Copy, Clone, PartialEq, Eq)]
116pub enum ChannelIntent {
117 Luma,
119 Chroma,
121 LumaChroma,
124 YuvFused,
128}
129
130impl ChannelIntent {
131 pub fn validate_for_source(self, layout: FrameLayout) -> Result<(), anyhow::Error> {
133 match self {
134 ChannelIntent::YuvFused if layout.subsampling != Subsampling::Yuv444 => {
135 anyhow::bail!(
136 "--channel-mode yuv requires a YUV444 source, got {:?}. Convert the input first, for example with `ffmpeg -pix_fmt yuv444p`",
137 layout.subsampling
138 );
139 },
140 _ => Ok(()),
141 }
142 }
143}
144
145#[derive(Debug, Clone)]
148pub struct PlaneOptions {
149 pub accelerators: Vec<Accelerator>,
150 pub device: Device,
151 pub intent: ChannelIntent,
152 pub mode: DenoisingMode,
153 pub algorithm: Algorithm,
156 pub luma_strength: Option<f32>,
160 pub chroma_strength: Option<f32>,
164 pub luma_lambda_ht: Option<f32>,
170 pub chroma_lambda_ht: Option<f32>,
176 pub luma_mismatch_scale: Option<f32>,
180 pub chroma_mismatch_scale: Option<f32>,
184}
185
186impl PlaneOptions {
187 fn algorithm_for(&self, channels: ChannelMode) -> Algorithm {
201 let per_plane = |luma, chroma| match channels {
202 ChannelMode::Luma => luma,
203 ChannelMode::Chroma => chroma,
204 ChannelMode::Yuv => None,
205 };
206
207 match self.algorithm {
208 Algorithm::Nl4d(nl4d) => Algorithm::Nl4d(Nl4dOptions {
209 lambda_ht: per_plane(self.luma_lambda_ht, self.chroma_lambda_ht).or(nl4d.lambda_ht),
213 mismatch_scale: per_plane(self.luma_mismatch_scale, self.chroma_mismatch_scale)
217 .unwrap_or(nl4d.mismatch_scale),
218 ..nl4d
219 }),
220 Algorithm::Nlmeans(nlm) => {
221 let strength = per_plane(self.luma_strength, self.chroma_strength);
222 Algorithm::Nlmeans(with_plane_strength(nlm, strength))
223 },
224 Algorithm::NlmeansHq(opts) => {
225 let strength = per_plane(self.luma_strength, self.chroma_strength);
226 Algorithm::NlmeansHq(NlmeansHqOptions {
227 nlm: with_plane_strength(opts.nlm, strength),
228 ..opts
229 })
230 },
231 }
232 }
233
234 fn denoiser_options(&self, channels: ChannelMode) -> DenoiserOptions {
235 DenoiserOptions::builder()
236 .channel_mode(channels)
237 .mode(self.mode)
238 .algorithm(self.algorithm_for(channels))
239 .build()
240 }
241}
242
243fn with_plane_strength(nlm: NlmeansOptions, strength: Option<f32>) -> NlmeansOptions {
246 match strength {
247 None => nlm,
248 Some(strength) => NlmeansOptions {
249 tuning: NlmTuning {
250 strength: Some(strength),
251 ..nlm.tuning
252 },
253 ..nlm
254 },
255 }
256}
257
258fn drop_leading<T>(queue: &mut VecDeque<T>, count: usize) {
260 for _ in 0..count.min(queue.len()) {
261 queue.pop_front();
262 }
263}
264
265pub fn push_needs_retry(result: Result<(), DenoiserError>) -> Result<bool, anyhow::Error> {
274 match result {
275 Ok(()) => Ok(false),
276 Err(DenoiserError::QueueFull) => Ok(true),
277 Err(other) => Err(other.into()),
278 }
279}
280
281pub struct PlanarDenoiser {
287 layout: FrameLayout,
288 luma: Option<Denoiser>,
289 chroma: Option<Denoiser>,
290 yuv: Option<Denoiser>,
293 luma_passthrough: VecDeque<Vec<u8>>,
298 chroma_passthrough: VecDeque<(Vec<u8>, Vec<u8>)>,
299 temporal_radius: u32,
302}
303
304impl PlanarDenoiser {
305 pub fn create(opts: &PlaneOptions, layout: FrameLayout) -> Result<Self, anyhow::Error> {
306 let (chroma_w, chroma_h) = layout.chroma_dims();
307
308 if chroma_w == 0 || chroma_h == 0 {
309 anyhow::bail!(
310 "frame dimensions {}x{} are too small for subsampling {:?}",
311 layout.width,
312 layout.height,
313 layout.subsampling
314 );
315 }
316
317 opts.intent.validate_for_source(layout)?;
318
319 let (denoise_luma, denoise_chroma, denoise_yuv) = match opts.intent {
320 ChannelIntent::Luma => (true, false, false),
321 ChannelIntent::Chroma => (false, true, false),
322 ChannelIntent::LumaChroma => (true, true, false),
323 ChannelIntent::YuvFused => (false, false, true),
324 };
325
326 let luma = denoise_luma
327 .then(|| {
328 Denoiser::create(
329 &opts.accelerators,
330 &opts.device,
331 layout.width,
332 layout.height,
333 opts.denoiser_options(ChannelMode::Luma),
334 )
335 })
336 .transpose()?;
337
338 let chroma = denoise_chroma
339 .then(|| {
340 Denoiser::create(
341 &opts.accelerators,
342 &opts.device,
343 chroma_w,
344 chroma_h,
345 opts.denoiser_options(ChannelMode::Chroma),
346 )
347 })
348 .transpose()?;
349
350 let yuv = denoise_yuv
351 .then(|| {
352 Denoiser::create(
353 &opts.accelerators,
354 &opts.device,
355 layout.width,
356 layout.height,
357 opts.denoiser_options(ChannelMode::Yuv),
358 )
359 })
360 .transpose()?;
361
362 let temporal_radius = match opts.mode {
363 DenoisingMode::Spacial => 0,
364 DenoisingMode::Temporal { radius } => radius,
365 };
366
367 Ok(Self {
368 layout,
369 luma,
370 chroma,
371 yuv,
372 luma_passthrough: VecDeque::new(),
373 chroma_passthrough: VecDeque::new(),
374 temporal_radius,
375 })
376 }
377
378 pub fn temporal_radius(&self) -> u32 {
380 self.temporal_radius
381 }
382
383 pub fn push(&mut self, planes: &Planes) -> Result<(), DenoiserError> {
411 self.push_with(planes, Denoiser::push_frame)
412 }
413
414 fn push_priming(&mut self, planes: &Planes) -> Result<(), DenoiserError> {
422 self.push_with(planes, Denoiser::push_frame_priming)
423 }
424
425 fn push_with(
431 &mut self,
432 planes: &Planes,
433 push_frame: fn(&mut Denoiser, &[f32]) -> Result<(), DenoiserError>,
434 ) -> Result<(), DenoiserError> {
435 if let Some(d) = self.yuv.as_mut() {
436 let buf = interleave_yuv_to_f32(&planes.y, &planes.u, &planes.v, self.layout.depth);
437 push_frame(d, &buf)?;
438 return Ok(());
439 }
440
441 if let Some(d) = self.luma.as_mut() {
442 let buf = plane_to_f32(&planes.y, self.layout.depth);
443 push_frame(d, &buf)?;
444 }
445
446 if let Some(d) = self.chroma.as_mut() {
447 let buf = interleave_uv_to_f32(&planes.u, &planes.v, self.layout.depth);
448 push_frame(d, &buf)?;
449 }
450
451 if self.luma.is_none() {
452 self.luma_passthrough.push_back(planes.y.clone());
453 }
454
455 if self.chroma.is_none() {
456 self.chroma_passthrough
457 .push_back((planes.u.clone(), planes.v.clone()));
458 }
459
460 Ok(())
461 }
462
463 pub fn recv(&mut self) -> Result<Option<Planes>, anyhow::Error> {
468 if let Some(d) = self.yuv.as_mut() {
469 return match d.recv_frame()? {
470 Some(packed) => Ok(Some(unpack_yuv_from_f32(
471 &packed,
472 self.layout.luma_pixels(),
473 self.layout.depth,
474 ))),
475 None => Ok(None),
476 };
477 }
478
479 let luma_out = self.luma.as_mut().map(|d| d.recv_frame()).transpose()?.flatten();
480
481 let chroma_out = self
482 .chroma
483 .as_mut()
484 .map(|d| d.recv_frame())
485 .transpose()?
486 .flatten();
487
488 let luma_passthrough = if self.luma.is_none() && chroma_out.is_some() {
492 self.luma_passthrough.pop_front()
493 } else {
494 None
495 };
496
497 let chroma_passthrough = if self.chroma.is_none() && luma_out.is_some() {
498 self.chroma_passthrough.pop_front()
499 } else {
500 None
501 };
502
503 if luma_out.is_none() && chroma_out.is_none() {
504 return Ok(None);
505 }
506
507 let planes = self.assemble(luma_out, chroma_out, luma_passthrough, chroma_passthrough);
508
509 Ok(Some(planes))
510 }
511
512 pub fn flush(&mut self, mut sink: impl FnMut(Planes)) -> Result<(), anyhow::Error> {
516 if let Some(d) = self.yuv.as_mut() {
517 let pixels = self.layout.luma_pixels();
518 let depth = self.layout.depth;
519 d.flush(|packed| sink(unpack_yuv_from_f32(&packed, pixels, depth)))?;
520 return Ok(());
521 }
522
523 let chroma_pixels = self.layout.chroma_pixels();
524
525 let mut luma_buf: Vec<Vec<f32>> = Vec::new();
526 let mut chroma_buf: Vec<Vec<f32>> = Vec::new();
527
528 if let Some(d) = self.luma.as_mut() {
529 d.flush(|v| luma_buf.push(v))?;
530 }
531
532 if let Some(d) = self.chroma.as_mut() {
533 d.flush(|v| chroma_buf.push(v))?;
534 }
535
536 let count = luma_buf.len().max(chroma_buf.len());
541
542 for i in 0..count {
543 let y = if let Some(buf) = luma_buf.get(i) {
544 f32_to_plane(buf, self.layout.depth)
545 } else if let Some(src) = self.luma_passthrough.pop_front() {
546 src
547 } else {
548 self.layout.black_luma_plane()
549 };
550
551 let (u, v) = if let Some(packed) = chroma_buf.get(i) {
552 unpack_uv_from_f32(packed, chroma_pixels, self.layout.depth)
553 } else if let Some((src_u, src_v)) = self.chroma_passthrough.pop_front() {
554 (src_u, src_v)
555 } else {
556 (
557 self.layout.neutral_chroma_plane(),
558 self.layout.neutral_chroma_plane(),
559 )
560 };
561
562 sink(Planes { y, u, v });
563 }
564
565 if !self.luma_passthrough.is_empty() || !self.chroma_passthrough.is_empty() {
566 tracing::warn!(
567 luma_remaining = self.luma_passthrough.len(),
568 chroma_remaining = self.chroma_passthrough.len(),
569 "passthrough queue not fully drained after flush",
570 );
571 self.luma_passthrough.clear();
572 self.chroma_passthrough.clear();
573 }
574
575 Ok(())
576 }
577
578 pub fn window_span(&self) -> WindowSpan {
585 self.yuv
586 .as_ref()
587 .or(self.luma.as_ref())
588 .or(self.chroma.as_ref())
589 .expect("PlanarDenoiser always keeps at least one Denoiser")
590 .window_span()
591 }
592
593 pub fn reseed(&mut self, window: &[Planes]) -> Result<Planes, anyhow::Error> {
625 let span = self.window_span();
626 let expected = span.frame_count();
627 if window.len() != expected {
628 anyhow::bail!("reseed needs a window of {expected} frames, got {}", window.len());
629 }
630
631 self.luma_passthrough.clear();
632 self.chroma_passthrough.clear();
633
634 for d in [self.yuv.as_mut(), self.luma.as_mut(), self.chroma.as_mut()]
635 .into_iter()
636 .flatten()
637 {
638 d.reset_stream();
639 }
640
641 let radius = self.temporal_radius as usize;
648 let priming_count = 2 * radius;
649 let (head, tail) = window.split_at(priming_count);
650 for planes in head {
651 self.push_priming(planes)?;
652 }
653
654 drop_leading(&mut self.luma_passthrough, radius);
674 drop_leading(&mut self.chroma_passthrough, radius);
675
676 let mut result = None;
677 for planes in tail {
678 self.push(planes)?;
679 if let Some(out) = self.recv()? {
680 result = Some(out);
681 }
682 }
683
684 result.ok_or_else(|| anyhow::anyhow!("a full window produced no frame, this is a bug"))
685 }
686
687 fn assemble(
688 &self,
689 luma: Option<Vec<f32>>,
690 chroma: Option<Vec<f32>>,
691 luma_passthrough: Option<Vec<u8>>,
692 chroma_passthrough: Option<(Vec<u8>, Vec<u8>)>,
693 ) -> Planes {
694 let chroma_pixels = self.layout.chroma_pixels();
695
696 let y = match (luma, luma_passthrough) {
697 (Some(v), _) => f32_to_plane(&v, self.layout.depth),
698 (None, Some(src)) => src,
699 (None, None) => self.layout.black_luma_plane(),
700 };
701
702 let (u, v) = match (chroma, chroma_passthrough) {
703 (Some(packed), _) => unpack_uv_from_f32(&packed, chroma_pixels, self.layout.depth),
704 (None, Some(src)) => src,
705 (None, None) => (
706 self.layout.neutral_chroma_plane(),
707 self.layout.neutral_chroma_plane(),
708 ),
709 };
710
711 Planes { y, u, v }
712 }
713}
714
715trait SampleCodec {
720 const BYTES: usize;
721
722 fn read(plane: &[u8], i: usize) -> u16;
723 fn write(plane: &mut [u8], i: usize, value: u16);
724}
725
726struct Narrow;
728
729impl SampleCodec for Narrow {
730 const BYTES: usize = 1;
731
732 #[inline(always)]
733 fn read(plane: &[u8], i: usize) -> u16 {
734 plane[i] as u16
735 }
736
737 #[inline(always)]
738 fn write(plane: &mut [u8], i: usize, value: u16) {
739 plane[i] = value as u8;
740 }
741}
742
743struct Wide;
745
746impl SampleCodec for Wide {
747 const BYTES: usize = 2;
748
749 #[inline(always)]
750 fn read(plane: &[u8], i: usize) -> u16 {
751 u16::from_le_bytes([plane[2 * i], plane[2 * i + 1]])
752 }
753
754 #[inline(always)]
755 fn write(plane: &mut [u8], i: usize, value: u16) {
756 plane[2 * i..2 * i + 2].copy_from_slice(&value.to_le_bytes());
757 }
758}
759
760#[inline(always)]
762fn quantise(v: f32, max: f32) -> u16 {
763 (v.clamp(0.0, 1.0) * max + 0.5) as u16
764}
765
766pub fn plane_to_f32(plane: &[u8], depth: Depth) -> Vec<f32> {
768 let max = depth.max_value();
769
770 fn run<C: SampleCodec>(plane: &[u8], max: f32) -> Vec<f32> {
771 let samples = plane.len() / C::BYTES;
772 (0..samples).map(|i| C::read(plane, i) as f32 / max).collect()
773 }
774
775 match depth.bytes_per_sample() {
776 1 => run::<Narrow>(plane, max),
777 _ => run::<Wide>(plane, max),
778 }
779}
780
781pub fn f32_to_plane(plane: &[f32], depth: Depth) -> Vec<u8> {
783 let max = depth.max_value();
784
785 fn run<C: SampleCodec>(plane: &[f32], max: f32) -> Vec<u8> {
786 let mut out = vec![0u8; plane.len() * C::BYTES];
787 for (i, &v) in plane.iter().enumerate() {
788 C::write(&mut out, i, quantise(v, max));
789 }
790 out
791 }
792
793 match depth.bytes_per_sample() {
794 1 => run::<Narrow>(plane, max),
795 _ => run::<Wide>(plane, max),
796 }
797}
798
799pub fn interleave_yuv_to_f32(y: &[u8], u: &[u8], v: &[u8], depth: Depth) -> Vec<f32> {
804 debug_assert_eq!(y.len(), u.len());
805 debug_assert_eq!(u.len(), v.len());
806
807 let max = depth.max_value();
808
809 fn run<C: SampleCodec>(y: &[u8], u: &[u8], v: &[u8], max: f32) -> Vec<f32> {
810 let pixels = y.len() / C::BYTES;
811 let mut out = Vec::with_capacity(pixels * 3);
812
813 for i in 0..pixels {
814 out.push(C::read(y, i) as f32 / max);
815 out.push(C::read(u, i) as f32 / max);
816 out.push(C::read(v, i) as f32 / max);
817 }
818
819 out
820 }
821
822 match depth.bytes_per_sample() {
823 1 => run::<Narrow>(y, u, v, max),
824 _ => run::<Wide>(y, u, v, max),
825 }
826}
827
828pub fn unpack_yuv_from_f32(packed: &[f32], pixels: usize, depth: Depth) -> Planes {
830 debug_assert_eq!(packed.len(), 3 * pixels);
831
832 let max = depth.max_value();
833
834 fn run<C: SampleCodec>(packed: &[f32], pixels: usize, max: f32) -> Planes {
835 let mut y = vec![0u8; pixels * C::BYTES];
836 let mut u = vec![0u8; pixels * C::BYTES];
837 let mut v = vec![0u8; pixels * C::BYTES];
838
839 for (i, chunk) in packed.as_chunks::<3>().0.iter().enumerate() {
840 C::write(&mut y, i, quantise(chunk[0], max));
841 C::write(&mut u, i, quantise(chunk[1], max));
842 C::write(&mut v, i, quantise(chunk[2], max));
843 }
844
845 Planes { y, u, v }
846 }
847
848 match depth.bytes_per_sample() {
849 1 => run::<Narrow>(packed, pixels, max),
850 _ => run::<Wide>(packed, pixels, max),
851 }
852}
853
854pub fn interleave_uv_to_f32(u: &[u8], v: &[u8], depth: Depth) -> Vec<f32> {
857 debug_assert_eq!(u.len(), v.len());
858
859 let max = depth.max_value();
860
861 fn run<C: SampleCodec>(u: &[u8], v: &[u8], max: f32) -> Vec<f32> {
862 let pixels = u.len() / C::BYTES;
863 let mut out = Vec::with_capacity(pixels * 2);
864
865 for i in 0..pixels {
866 out.push(C::read(u, i) as f32 / max);
867 out.push(C::read(v, i) as f32 / max);
868 }
869
870 out
871 }
872
873 match depth.bytes_per_sample() {
874 1 => run::<Narrow>(u, v, max),
875 _ => run::<Wide>(u, v, max),
876 }
877}
878
879pub fn unpack_uv_from_f32(packed: &[f32], chroma_pixels: usize, depth: Depth) -> (Vec<u8>, Vec<u8>) {
881 debug_assert_eq!(packed.len(), 2 * chroma_pixels);
882
883 let max = depth.max_value();
884
885 fn run<C: SampleCodec>(packed: &[f32], chroma_pixels: usize, max: f32) -> (Vec<u8>, Vec<u8>) {
886 let mut u = vec![0u8; chroma_pixels * C::BYTES];
887 let mut v = vec![0u8; chroma_pixels * C::BYTES];
888
889 for (i, chunk) in packed.as_chunks::<2>().0.iter().enumerate() {
890 C::write(&mut u, i, quantise(chunk[0], max));
891 C::write(&mut v, i, quantise(chunk[1], max));
892 }
893
894 (u, v)
895 }
896
897 match depth.bytes_per_sample() {
898 1 => run::<Narrow>(packed, chroma_pixels, max),
899 _ => run::<Wide>(packed, chroma_pixels, max),
900 }
901}
902
903#[cfg(test)]
904mod converter_tests {
905 use super::*;
906
907 fn wire(samples: &[u16], depth: Depth) -> Vec<u8> {
910 match depth.bytes_per_sample() {
911 1 => samples.iter().map(|&s| s as u8).collect(),
912 _ => samples.iter().flat_map(|&s| s.to_le_bytes()).collect(),
913 }
914 }
915
916 #[test]
917 fn plane_round_trips_boundary_codes_at_every_depth() {
918 for depth in [Depth::Eight, Depth::Ten, Depth::Twelve] {
919 let max = depth.max_value() as u16;
920 let samples: Vec<u16> = vec![0, 1, 16, 64, 235, max / 2, max - 1, max]
921 .into_iter()
922 .filter(|&s| s <= max)
923 .collect();
924
925 let bytes = wire(&samples, depth);
926 let restored = f32_to_plane(&plane_to_f32(&bytes, depth), depth);
927
928 assert_eq!(restored, bytes, "plane round trip failed at {depth:?}");
929 }
930 }
931
932 #[test]
935 fn high_depth_samples_are_little_endian() {
936 let bytes = wire(&[1023, 0, 512], Depth::Ten);
938 assert_eq!(bytes, vec![0xFF, 0x03, 0x00, 0x00, 0x00, 0x02]);
939
940 let f = plane_to_f32(&bytes, Depth::Ten);
941 assert!(
942 (f[0] - 1.0).abs() < 1e-6,
943 "0x03FF should normalize to 1.0, got {}",
944 f[0]
945 );
946 assert_eq!(f[1], 0.0);
947 }
948
949 #[test]
950 fn uv_interleave_round_trips_at_every_depth() {
951 for depth in [Depth::Eight, Depth::Ten, Depth::Twelve] {
952 let max = depth.max_value() as u16;
953 let u_samples = vec![0, max / 4, max];
954 let v_samples = vec![max, max / 2, 1];
955
956 let u_bytes = wire(&u_samples, depth);
957 let v_bytes = wire(&v_samples, depth);
958
959 let packed = interleave_uv_to_f32(&u_bytes, &v_bytes, depth);
960 assert_eq!(packed.len(), 6, "packed UV length wrong at {depth:?}");
961
962 let (ru, rv) = unpack_uv_from_f32(&packed, 3, depth);
963 assert_eq!(ru, u_bytes, "U round trip failed at {depth:?}");
964 assert_eq!(rv, v_bytes, "V round trip failed at {depth:?}");
965 }
966 }
967
968 #[test]
969 fn yuv_interleave_round_trips_at_every_depth() {
970 for depth in [Depth::Eight, Depth::Ten, Depth::Twelve] {
971 let max = depth.max_value() as u16;
972 let y_samples = vec![0, max / 3, max];
973 let u_samples = vec![max, 0, max / 2];
974 let v_samples = vec![max / 4, max, 0];
975
976 let y_bytes = wire(&y_samples, depth);
977 let u_bytes = wire(&u_samples, depth);
978 let v_bytes = wire(&v_samples, depth);
979
980 let packed = interleave_yuv_to_f32(&y_bytes, &u_bytes, &v_bytes, depth);
981 assert_eq!(packed.len(), 9, "packed YUV length wrong at {depth:?}");
982
983 let out = unpack_yuv_from_f32(&packed, 3, depth);
984 assert_eq!(out.y, y_bytes, "Y round trip failed at {depth:?}");
985 assert_eq!(out.u, u_bytes, "U round trip failed at {depth:?}");
986 assert_eq!(out.v, v_bytes, "V round trip failed at {depth:?}");
987 }
988 }
989
990 #[test]
991 fn quantise_matches_the_clamping_form_including_nan() {
992 fn reference(v: f32, max: f32) -> u16 {
993 (v.clamp(0.0, 1.0) * max + 0.5) as u16
994 }
995
996 let max = 1023.0;
997 let cases = [
998 -1.0,
999 -0.001,
1000 0.0,
1001 0.5,
1002 0.999,
1003 1.0,
1004 1.001,
1005 2.0,
1006 f32::NAN,
1007 f32::INFINITY,
1008 f32::NEG_INFINITY,
1009 ];
1010
1011 for v in cases {
1012 assert_eq!(quantise(v, max), reference(v, max), "mismatch at {v}");
1013 }
1014 }
1015
1016 #[test]
1031 fn limited_range_codes_agree_across_depths() {
1032 const TOL: f32 = 1.0 / 255.0;
1034
1035 let eight = plane_to_f32(&wire(&[16, 235], Depth::Eight), Depth::Eight);
1036 let ten = plane_to_f32(&wire(&[64, 940], Depth::Ten), Depth::Ten);
1037
1038 for (a, b) in eight.iter().zip(ten.iter()) {
1039 assert!((a - b).abs() < TOL, "8-bit {a} vs 10-bit {b}");
1040 }
1041 }
1042}
1043
1044#[cfg(test)]
1045mod cli_options_tests {
1046 use super::*;
1047 use crate::nlmeans::NlmParams;
1048
1049 fn base_opts(
1055 mode: DenoisingMode,
1056 algorithm: Algorithm,
1057 luma_strength: Option<f32>,
1058 chroma_strength: Option<f32>,
1059 ) -> PlaneOptions {
1060 PlaneOptions {
1061 accelerators: vec![],
1062 device: Device::Default,
1063 intent: ChannelIntent::LumaChroma,
1064 mode,
1065 algorithm,
1066 luma_strength,
1067 chroma_strength,
1068 luma_lambda_ht: None,
1069 chroma_lambda_ht: None,
1070 luma_mismatch_scale: None,
1071 chroma_mismatch_scale: None,
1072 }
1073 }
1074
1075 #[test]
1076 fn luma_strength_alone_overrides_only_the_luma_plane() {
1077 let opts = base_opts(DenoisingMode::Spacial, Algorithm::default(), Some(0.7), None);
1078
1079 let luma = expect_nlmeans(opts.denoiser_options(ChannelMode::Luma).algorithm);
1080 let chroma = expect_nlmeans(opts.denoiser_options(ChannelMode::Chroma).algorithm);
1081
1082 assert!(
1083 matches!(luma.tuning.strength, Some(s) if (s - 0.7).abs() < f32::EPSILON),
1084 "expected luma tuning.strength = Some(0.7), got {:?}",
1085 luma.tuning.strength
1086 );
1087 assert_eq!(
1088 chroma.tuning.strength, None,
1089 "chroma plane should carry no override so the table default applies"
1090 );
1091 }
1092
1093 #[test]
1094 fn both_per_plane_strengths_set_independently() {
1095 let opts = base_opts(DenoisingMode::Spacial, Algorithm::default(), Some(0.7), Some(0.3));
1096
1097 let luma = expect_nlmeans(opts.denoiser_options(ChannelMode::Luma).algorithm);
1098 let chroma = expect_nlmeans(opts.denoiser_options(ChannelMode::Chroma).algorithm);
1099
1100 assert!(
1101 matches!(luma.tuning.strength, Some(s) if (s - 0.7).abs() < f32::EPSILON),
1102 "expected luma tuning.strength = Some(0.7), got {:?}",
1103 luma.tuning.strength
1104 );
1105 assert!(
1106 matches!(chroma.tuning.strength, Some(s) if (s - 0.3).abs() < f32::EPSILON),
1107 "expected chroma tuning.strength = Some(0.3), got {:?}",
1108 chroma.tuning.strength
1109 );
1110 }
1111
1112 #[test]
1113 fn no_overrides_hq_resolves_through_to_nlm_params_to_the_measured_tables() {
1114 let opts = base_opts(
1117 DenoisingMode::Temporal { radius: 4 },
1118 Algorithm::NlmeansHq(NlmeansHqOptions::default()),
1119 None,
1120 None,
1121 );
1122
1123 let luma_params: NlmParams = opts.denoiser_options(ChannelMode::Luma).to_nlm_params();
1124 let chroma_params: NlmParams = opts.denoiser_options(ChannelMode::Chroma).to_nlm_params();
1125
1126 assert!(
1127 (luma_params.strength - 0.35).abs() < f32::EPSILON,
1128 "expected luma strength 0.35 at r4, got {}",
1129 luma_params.strength
1130 );
1131 assert!(
1132 (chroma_params.strength - 0.70).abs() < f32::EPSILON,
1133 "expected chroma strength 0.70 at r4, got {}",
1134 chroma_params.strength
1135 );
1136 }
1137
1138 fn nl4d_opts(luma_lambda_ht: Option<f32>, chroma_lambda_ht: Option<f32>) -> PlaneOptions {
1142 PlaneOptions {
1143 accelerators: vec![],
1144 device: Device::Default,
1145 intent: ChannelIntent::LumaChroma,
1146 mode: DenoisingMode::Temporal { radius: 2 },
1147 algorithm: Algorithm::Nl4d(Nl4dOptions::default()),
1148 luma_strength: None,
1149 chroma_strength: None,
1150 luma_lambda_ht,
1151 chroma_lambda_ht,
1152 luma_mismatch_scale: None,
1153 chroma_mismatch_scale: None,
1154 }
1155 }
1156
1157 fn expect_nlmeans(algorithm: Algorithm) -> NlmeansOptions {
1160 match algorithm {
1161 Algorithm::Nlmeans(n) => n,
1162 other => panic!("expected Algorithm::Nlmeans, got {other:?}"),
1163 }
1164 }
1165
1166 fn expect_nl4d(algorithm: Algorithm) -> Nl4dOptions {
1169 match algorithm {
1170 Algorithm::Nl4d(n) => n,
1171 other => panic!("expected Algorithm::Nl4d, got {other:?}"),
1172 }
1173 }
1174
1175 fn nl4d_mismatch_opts(
1178 shared: f32,
1179 luma_mismatch_scale: Option<f32>,
1180 chroma_mismatch_scale: Option<f32>,
1181 ) -> PlaneOptions {
1182 PlaneOptions {
1183 algorithm: Algorithm::Nl4d(Nl4dOptions {
1184 mismatch_scale: shared,
1185 ..Nl4dOptions::default()
1186 }),
1187 luma_mismatch_scale,
1188 chroma_mismatch_scale,
1189 ..nl4d_opts(None, None)
1190 }
1191 }
1192
1193 #[test]
1198 fn a_per_plane_mismatch_scale_overrides_only_its_own_instance_for_nl4d() {
1199 let luma_only = nl4d_mismatch_opts(2.0, Some(8.0), None);
1200 let luma = expect_nl4d(luma_only.algorithm_for(ChannelMode::Luma));
1201 let chroma = expect_nl4d(luma_only.algorithm_for(ChannelMode::Chroma));
1202 assert!((luma.mismatch_scale - 8.0).abs() < f32::EPSILON);
1203 assert!(
1204 (chroma.mismatch_scale - 2.0).abs() < f32::EPSILON,
1205 "chroma should keep the shared value, got {}",
1206 chroma.mismatch_scale
1207 );
1208
1209 let chroma_only = nl4d_mismatch_opts(2.0, None, Some(8.0));
1210 let luma = expect_nl4d(chroma_only.algorithm_for(ChannelMode::Luma));
1211 let chroma = expect_nl4d(chroma_only.algorithm_for(ChannelMode::Chroma));
1212 assert!((chroma.mismatch_scale - 8.0).abs() < f32::EPSILON);
1213 assert!(
1214 (luma.mismatch_scale - 2.0).abs() < f32::EPSILON,
1215 "luma should keep the shared value, got {}",
1216 luma.mismatch_scale
1217 );
1218 }
1219
1220 #[test]
1223 fn a_yuv_instance_ignores_both_per_plane_mismatch_scales() {
1224 let opts = nl4d_mismatch_opts(2.0, Some(8.0), Some(4.0));
1225 let yuv = expect_nl4d(opts.algorithm_for(ChannelMode::Yuv));
1226
1227 assert!((yuv.mismatch_scale - 2.0).abs() < f32::EPSILON);
1228 }
1229
1230 #[test]
1235 fn luma_lambda_ht_alone_overrides_only_the_luma_instance_for_nl4d() {
1236 let opts = nl4d_opts(Some(4.0), None);
1237
1238 let luma = expect_nl4d(opts.algorithm_for(ChannelMode::Luma));
1239 let chroma = expect_nl4d(opts.algorithm_for(ChannelMode::Chroma));
1240
1241 assert!((luma.lambda_ht.unwrap() - 4.0).abs() < f32::EPSILON);
1242 assert_eq!(
1243 chroma.lambda_ht,
1244 Nl4dOptions::default().lambda_ht,
1245 "chroma should stay unresolved here (None), deferred to its own per-plane \
1246 default at construction, got {:?}",
1247 chroma.lambda_ht
1248 );
1249 }
1250
1251 #[test]
1252 fn chroma_lambda_ht_alone_overrides_only_the_chroma_instance_for_nl4d() {
1253 let opts = nl4d_opts(None, Some(4.0));
1254
1255 let luma = expect_nl4d(opts.algorithm_for(ChannelMode::Luma));
1256 let chroma = expect_nl4d(opts.algorithm_for(ChannelMode::Chroma));
1257
1258 assert_eq!(
1259 luma.lambda_ht,
1260 Nl4dOptions::default().lambda_ht,
1261 "luma should stay unresolved here (None), deferred to its own per-plane \
1262 default at construction, got {:?}",
1263 luma.lambda_ht
1264 );
1265 assert!((chroma.lambda_ht.unwrap() - 4.0).abs() < f32::EPSILON);
1266 }
1267
1268 #[test]
1269 fn both_planes_lambda_ht_set_independently_for_nl4d() {
1270 let opts = nl4d_opts(Some(2.0), Some(3.5));
1271
1272 let luma = expect_nl4d(opts.algorithm_for(ChannelMode::Luma));
1273 let chroma = expect_nl4d(opts.algorithm_for(ChannelMode::Chroma));
1274
1275 assert!((luma.lambda_ht.unwrap() - 2.0).abs() < f32::EPSILON);
1276 assert!((chroma.lambda_ht.unwrap() - 3.5).abs() < f32::EPSILON);
1277
1278 assert_eq!(luma.refine, chroma.refine);
1281 assert_eq!(luma.spatial_radius, chroma.spatial_radius);
1282 assert!((luma.c_min - chroma.c_min).abs() < f32::EPSILON);
1283 }
1284
1285 #[test]
1286 fn unset_nl4d_overrides_resolve_to_different_lambda_ht_per_plane_end_to_end() {
1287 let opts = nl4d_opts(None, None);
1288
1289 let luma = expect_nl4d(opts.algorithm_for(ChannelMode::Luma));
1290 let chroma = expect_nl4d(opts.algorithm_for(ChannelMode::Chroma));
1291
1292 assert_eq!(luma.lambda_ht, None);
1295 assert_eq!(chroma.lambda_ht, None);
1296
1297 let luma_default = crate::nl4d_default_lambda_ht(ChannelMode::Luma);
1303 let chroma_default = crate::nl4d_default_lambda_ht(ChannelMode::Chroma);
1304 assert!((luma_default - 5.3).abs() < f32::EPSILON);
1305 assert!((chroma_default - 4.2).abs() < f32::EPSILON);
1306 assert!((chroma_default - luma_default).abs() > f32::EPSILON);
1307 }
1308}
1309
1310#[cfg(feature = "vulkan")]
1314#[cfg(test)]
1315mod passthrough_retry_tests {
1316 use super::*;
1317 use crate::accelerate::Accelerator;
1318 use crate::{Algorithm, DenoisingMode};
1319
1320 fn chroma_only_opts() -> PlaneOptions {
1326 PlaneOptions {
1327 accelerators: vec![Accelerator::Vulkan],
1328 device: Device::Default,
1329 intent: ChannelIntent::Chroma,
1330 mode: DenoisingMode::Spacial,
1331 algorithm: Algorithm::default(),
1332 luma_strength: None,
1333 chroma_strength: None,
1334 luma_lambda_ht: None,
1335 chroma_lambda_ht: None,
1336 luma_mismatch_scale: None,
1337 chroma_mismatch_scale: None,
1338 }
1339 }
1340
1341 fn fake_planes(layout: FrameLayout) -> Planes {
1342 Planes {
1343 y: fill_plane(layout.luma_pixels(), layout.depth.neutral_chroma(), layout.depth),
1344 u: layout.neutral_chroma_plane(),
1345 v: layout.neutral_chroma_plane(),
1346 }
1347 }
1348
1349 #[test]
1350 fn queue_full_retry_does_not_double_queue_the_passthrough_plane() {
1351 let layout = FrameLayout {
1352 width: 16,
1353 height: 16,
1354 subsampling: Subsampling::Yuv420,
1355 depth: Depth::Eight,
1356 };
1357 let mut wd =
1358 PlanarDenoiser::create(&chroma_only_opts(), layout).expect("denoiser construction failed");
1359 let planes = fake_planes(layout);
1360
1361 wd.push(&planes).expect("first push should land");
1365 wd.push(&planes).expect("second push should land");
1366
1367 let err = wd.push(&planes).expect_err("expected QueueFull");
1369 assert!(
1370 matches!(err, DenoiserError::QueueFull),
1371 "expected QueueFull, got {err:?}"
1372 );
1373
1374 wd.recv().expect("recv after drain failed");
1377 wd.push(&planes).expect("retry push should land after drain");
1378
1379 assert_eq!(
1385 wd.luma_passthrough.len(),
1386 2,
1387 "expected exactly one passthrough entry per chroma frame actually accepted, got {}",
1388 wd.luma_passthrough.len()
1389 );
1390 }
1391}
1392
1393#[cfg(feature = "vulkan")]
1397#[cfg(test)]
1398mod lumachroma_lockstep_tests {
1399 use super::*;
1400 use crate::accelerate::Accelerator;
1401 use crate::{Algorithm, DenoisingMode};
1402
1403 fn luma_chroma_opts() -> PlaneOptions {
1410 PlaneOptions {
1411 accelerators: vec![Accelerator::Vulkan],
1412 device: Device::Default,
1413 intent: ChannelIntent::LumaChroma,
1414 mode: DenoisingMode::Spacial,
1415 algorithm: Algorithm::default(),
1416 luma_strength: None,
1417 chroma_strength: None,
1418 luma_lambda_ht: None,
1419 chroma_lambda_ht: None,
1420 luma_mismatch_scale: None,
1421 chroma_mismatch_scale: None,
1422 }
1423 }
1424
1425 fn marked_planes(layout: FrameLayout, idx: u8) -> Planes {
1431 let chroma_pixels = layout.chroma_pixels();
1432 let y_val = 10 + idx;
1433 let uv_val = 200 - idx;
1434
1435 Planes {
1436 y: fill_plane(layout.luma_pixels(), y_val as u16, layout.depth),
1437 u: fill_plane(chroma_pixels, uv_val as u16, layout.depth),
1438 v: fill_plane(chroma_pixels, uv_val as u16, layout.depth),
1439 }
1440 }
1441
1442 #[test]
1443 fn queue_full_retries_never_desync_luma_and_chroma() {
1444 let layout = FrameLayout {
1445 width: 16,
1446 height: 16,
1447 subsampling: Subsampling::Yuv420,
1448 depth: Depth::Eight,
1449 };
1450 let mut wd =
1451 PlanarDenoiser::create(&luma_chroma_opts(), layout).expect("denoiser construction failed");
1452
1453 const N: u8 = 6;
1456 let mut outputs: Vec<Planes> = Vec::new();
1457
1458 for idx in 0..N {
1459 let planes = marked_planes(layout, idx);
1460
1461 if push_needs_retry(wd.push(&planes)).expect("push_needs_retry") {
1464 if let Some(out) = wd.recv().expect("recv failed") {
1465 outputs.push(out);
1466 }
1467
1468 wd.push(&planes).expect("retry push should land after drain");
1469 }
1470 }
1471
1472 wd.flush(|out| outputs.push(out)).expect("flush failed");
1473
1474 assert_eq!(
1475 outputs.len(),
1476 N as usize,
1477 "expected exactly one output frame per input frame, got {}",
1478 outputs.len()
1479 );
1480
1481 for out in &outputs {
1482 let y_val = out.y[0];
1483 let uv_val = out.u[0];
1484 let idx_from_y = y_val - 10;
1485 let idx_from_uv = 200 - uv_val;
1486
1487 assert_eq!(
1488 idx_from_y, idx_from_uv,
1489 "luma marker {y_val} (frame {idx_from_y}) and chroma marker {uv_val} \
1490 (frame {idx_from_uv}) disagree, so the luma and chroma pushes have drifted apart"
1491 );
1492 }
1493 }
1494}
1495
1496#[cfg(test)]
1497mod push_needs_retry_tests {
1498 use super::*;
1499
1500 #[test]
1501 fn ok_means_no_retry() {
1502 let outcome = push_needs_retry(Ok(())).expect("Ok(()) must not itself error");
1503 assert!(!outcome, "a landed push must not ask the caller to retry");
1504 }
1505
1506 #[test]
1507 fn queue_full_signals_retry() {
1508 let outcome =
1509 push_needs_retry(Err(DenoiserError::QueueFull)).expect("QueueFull must not itself error");
1510 assert!(outcome, "QueueFull must still trigger the retry-after-drain path");
1511 }
1512
1513 #[test]
1514 fn non_queue_full_errors_propagate_instead_of_being_swallowed() {
1515 let synthetic = DenoiserError::Other(anyhow::anyhow!("synthetic readback failure"));
1516
1517 let outcome = push_needs_retry(Err(synthetic));
1518
1519 assert!(
1520 outcome.is_err(),
1521 "a non-QueueFull push error must propagate instead of being silently treated as success"
1522 );
1523 }
1524}
1525
1526#[cfg(test)]
1527mod layout_tests {
1528 use super::*;
1529
1530 fn layout(depth: Depth) -> FrameLayout {
1531 FrameLayout {
1532 width: 4,
1533 height: 4,
1534 subsampling: Subsampling::Yuv420,
1535 depth,
1536 }
1537 }
1538
1539 #[test]
1540 fn byte_lengths_scale_with_depth() {
1541 assert_eq!(layout(Depth::Eight).luma_bytes(), 16);
1542 assert_eq!(layout(Depth::Ten).luma_bytes(), 32);
1543 assert_eq!(layout(Depth::Eight).chroma_bytes(), 4);
1544 assert_eq!(layout(Depth::Ten).chroma_bytes(), 8);
1545 }
1546
1547 #[test]
1548 fn neutral_chroma_fill_is_correct_at_each_depth() {
1549 let eight = layout(Depth::Eight).neutral_chroma_plane();
1550 assert_eq!(eight, vec![128u8; 4]);
1551
1552 let ten = layout(Depth::Ten).neutral_chroma_plane();
1554 assert_eq!(ten, vec![0x00, 0x02, 0x00, 0x02, 0x00, 0x02, 0x00, 0x02]);
1555
1556 let twelve = layout(Depth::Twelve).neutral_chroma_plane();
1558 assert_eq!(twelve.len(), 8);
1559 assert_eq!(&twelve[0..2], &[0x00, 0x08]);
1560 }
1561
1562 #[test]
1563 fn black_luma_fill_is_zero_at_the_right_length() {
1564 assert_eq!(layout(Depth::Eight).black_luma_plane(), vec![0u8; 16]);
1565 assert_eq!(layout(Depth::Ten).black_luma_plane(), vec![0u8; 32]);
1566 }
1567}
1568
1569#[cfg(test)]
1570mod tests;