Skip to main content

nuviz_cli/commands/
diff.rs

1use std::path::Path;
2
3use anyhow::{bail, Result};
4
5use crate::cli::DiffArgs;
6use crate::data::experiment::discover_experiments;
7use crate::data::images::discover_images;
8use crate::terminal::capability::detect_capabilities;
9use crate::terminal::heatmap::generate_error_heatmap;
10use crate::terminal::render;
11
12pub fn run(args: DiffArgs, base_dir: &Path) -> Result<()> {
13    let experiments = discover_experiments(base_dir);
14
15    let exp_a = experiments
16        .iter()
17        .find(|e| {
18            e.name == args.experiment_a
19                || e.dir.file_name().and_then(|n| n.to_str()) == Some(&args.experiment_a)
20        })
21        .ok_or_else(|| anyhow::anyhow!("Experiment '{}' not found", args.experiment_a))?;
22
23    let exp_b = experiments
24        .iter()
25        .find(|e| {
26            e.name == args.experiment_b
27                || e.dir.file_name().and_then(|n| n.to_str()) == Some(&args.experiment_b)
28        })
29        .ok_or_else(|| anyhow::anyhow!("Experiment '{}' not found", args.experiment_b))?;
30
31    let caps = detect_capabilities();
32    let tag = &args.tag;
33
34    // Find matching image pairs
35    let images_a = discover_images(&exp_a.dir);
36    let images_b = discover_images(&exp_b.dir);
37
38    let images_a_filtered: Vec<_> = images_a.iter().filter(|e| e.tag == *tag).collect();
39    let images_b_filtered: Vec<_> = images_b.iter().filter(|e| e.tag == *tag).collect();
40
41    // Determine which step to compare
42    let step = if let Some(s) = args.step {
43        s
44    } else {
45        // Find the latest common step
46        let steps_a: std::collections::HashSet<u64> =
47            images_a_filtered.iter().map(|e| e.step).collect();
48        let steps_b: std::collections::HashSet<u64> =
49            images_b_filtered.iter().map(|e| e.step).collect();
50        let common: Vec<u64> = steps_a.intersection(&steps_b).copied().collect();
51
52        if common.is_empty() {
53            bail!(
54                "No common steps found for tag '{}' between '{}' and '{}'",
55                tag,
56                args.experiment_a,
57                args.experiment_b
58            );
59        }
60        *common.iter().max().unwrap()
61    };
62
63    let img_entry_a = images_a_filtered
64        .iter()
65        .find(|e| e.step == step)
66        .ok_or_else(|| anyhow::anyhow!("No image at step {} in '{}'", step, args.experiment_a))?;
67
68    let img_entry_b = images_b_filtered
69        .iter()
70        .find(|e| e.step == step)
71        .ok_or_else(|| anyhow::anyhow!("No image at step {} in '{}'", step, args.experiment_b))?;
72
73    println!(
74        "Comparing step={step} tag={tag}: {} vs {}",
75        args.experiment_a, args.experiment_b
76    );
77
78    let (max_w, max_h) = render::get_terminal_pixel_size();
79
80    if args.heatmap {
81        // Generate and display error heatmap
82        let img_a = image::open(&img_entry_a.path)?;
83        let img_b = image::open(&img_entry_b.path)?;
84        let result = generate_error_heatmap(&img_a, &img_b);
85
86        render::render_dynamic_image(&result.image, &caps, max_w, max_h.saturating_sub(128))?;
87
88        println!();
89        println!("Error Analysis:");
90        println!("  MAE:       {:.4}", result.mae);
91        println!("  Max Error: {}", result.max_error);
92        if result.psnr.is_infinite() {
93            println!("  PSNR:      ∞ (identical images)");
94        } else {
95            println!("  PSNR:      {:.2} dB", result.psnr);
96        }
97    } else {
98        // Side-by-side comparison
99        render::render_image_pair(
100            &img_entry_a.path,
101            &img_entry_b.path,
102            &caps,
103            max_w,
104            max_h.saturating_sub(128),
105        )?;
106
107        // Also compute and display metrics
108        let img_a = image::open(&img_entry_a.path)?;
109        let img_b = image::open(&img_entry_b.path)?;
110        let result = generate_error_heatmap(&img_a, &img_b);
111
112        println!();
113        println!(
114            "  {} (left) vs {} (right)",
115            args.experiment_a, args.experiment_b
116        );
117        if result.psnr.is_infinite() {
118            println!("  PSNR: ∞ | MAE: {:.4}", result.mae);
119        } else {
120            println!("  PSNR: {:.2} dB | MAE: {:.4}", result.psnr, result.mae);
121        }
122    }
123
124    Ok(())
125}