DarkRoom/core/dr-pano/examples/align.rs

//! Align real frames from their embedded previews and draw the result.
//!
//! ```sh
//! cargo run -p dr-pano --example align --release -- fixtures/pano/2025-08-05/*.CR2
//! cargo run -p dr-pano --example align --release -- out-prefix frame1.CR2 frame2.CR2 …
//! ```
//!
//! The point of looking rather than asserting: a rotation solve that is
//! numerically converged and geometrically wrong — a mirrored axis, a
//! transposed homography, an orientation applied the wrong way — produces
//! perfectly plausible residuals and a picture that is obviously broken.
//! This writes `<prefix>-cyl.ppm`: every frame's preview warped onto a
//! cylinder and averaged where they overlap, at a size that fits on a
//! screen. Ghosting in the overlaps is the alignment error, made visible.
//!
//! Previews, not RAW: the alignment runs on proxies in the application too
//! (FR-MRG-7), and a camera's embedded JPEG is a proxy the decoder already
//! extracts in milliseconds. What is different from the real path is only
//! that the pixels are the camera's rendering rather than ours, which the
//! geometry does not care about.

use std::path::PathBuf;
use std::time::Instant;

use dr_pano::bundle::Cameras;
use dr_pano::{align, xfeat::XFeat, AlignOptions, Gray, Projection};

fn main() {
    env_logger::init();
    let mut args: Vec<String> = std::env::args().skip(1).collect();
    if args.is_empty() {
        eprintln!("usage: align [out-prefix] <frame>...");
        std::process::exit(2);
    }
    let prefix = if args[0].ends_with(".CR2") || args[0].ends_with(".dng") || args[0].ends_with(".jpg") {
        "align".to_string()
    } else {
        args.remove(0)
    };
    let paths: Vec<PathBuf> = args.iter().map(PathBuf::from).collect();

    // Previews, oriented, at proxy size.
    let t = Instant::now();
    let mut proxies: Vec<Gray> = Vec::new();
    for p in &paths {
        let bytes = std::fs::read(p).expect("read");
        let preview = dr_decode::extract_preview(&bytes, dr_decode::PreviewSize::Full)
            .expect("embedded preview");
        let orientation = dr_decode::orientation(&bytes[..bytes.len().min(dr_decode::HEADER_BYTES as usize)])
            .unwrap_or(dr_types::Orientation::NORMAL);
        let tag = match orientation.quarter_turns {
            1 => 6,
            2 => 3,
            3 => 8,
            _ => 1,
        };
        let gray = Gray::from_rgba8(&preview.rgba, preview.width as usize, preview.height as usize)
            .oriented(tag);
        let (fitted, _) = gray.fitted(dr_pano::xfeat::INPUT_LONG_EDGE, dr_pano::xfeat::INPUT_LONG_EDGE);
        println!(
            "{:<14} preview {}×{} orientation {} → proxy {}×{}",
            p.file_name().unwrap().to_string_lossy(),
            preview.width,
            preview.height,
            tag,
            fitted.width,
            fitted.height
        );
        proxies.push(fitted);
    }
    println!("previews in {:?}", t.elapsed());

    // Keypoints.
    let t = Instant::now();
    let mut detector = XFeat::embedded().expect("model");
    let features: Vec<_> = proxies
        .iter()
        .map(|g| detector.detect(g).expect("detect"))
        .collect();
    for (i, f) in features.iter().enumerate() {
        println!("frame {i}: {} keypoints", f.len());
    }
    println!("detection in {:?} ({:?} per frame)", t.elapsed(), t.elapsed() / proxies.len() as u32);

    // Alignment.
    let t = Instant::now();
    let opts = AlignOptions::default();
    let alignment = align(&features, &opts).expect("align");
    println!("alignment in {:?}", t.elapsed());
    println!(
        "focal {:.1} px, long edge {} px ({:.1} mm on full frame), rms {:.3} px",
        alignment.focal,
        proxies[0].width.max(proxies[0].height),
        alignment.focal * 36.0 / proxies[0].width.max(proxies[0].height) as f64,
        alignment.rms_px
    );
    for l in &alignment.links {
        println!("  link {}–{}: {} inliers of {} matches", l.i, l.j, l.inliers, l.matches);
    }
    for (k, why) in &alignment.unaligned {
        println!("  UNALIGNED frame {k}: {why}");
    }
    let root = alignment
        .rotations
        .iter()
        .position(|r| *r == Some(dr_pano::linalg::Mat3::IDENTITY))
        .unwrap_or(0);
    for (k, r) in alignment.rotations.iter().enumerate() {
        if let Some(r) = r {
            // Yaw about y, pitch about x, roll about z, from the matrix's
            // columns — enough to read a sweep by eye.
            let yaw = r.0[0][2].atan2(r.0[2][2]).to_degrees();
            let pitch = (-r.0[1][2]).asin().to_degrees();
            let roll = r.0[1][0].atan2(r.0[1][1]).to_degrees();
            println!(
                "  frame {k}: yaw {yaw:7.2}°  pitch {pitch:6.2}°  roll {roll:6.2}°{}",
                if k == root { "  (reference)" } else { "" }
            );
        }
    }

    if !alignment.is_complete() {
        eprintln!("not drawing: the set is not fully aligned");
        std::process::exit(1);
    }

    // Draw: a cylinder, averaged where frames overlap.
    let t = Instant::now();
    let cameras: Cameras = alignment.cameras();
    let (fw, fh) = (proxies[0].width as f64, proxies[0].height as f64);
    let scale = alignment.focal;
    let bounds = dr_pano::projection::bounds(Projection::Cylindrical, scale, &cameras, (fw, fh))
        .expect("bounds");
    // Fit to 3000 px wide.
    let out_w = 3000usize;
    let px = bounds.width() / out_w as f64;
    let out_h = (bounds.height() / px).ceil() as usize;
    let mut sum = vec![0.0f32; out_w * out_h];
    let mut count = vec![0u16; out_w * out_h];
    for oy in 0..out_h {
        for ox in 0..out_w {
            let u = bounds.min_u + (ox as f64 + 0.5) * px;
            let v = bounds.min_v + (oy as f64 + 0.5) * px;
            let d = Projection::Cylindrical.to_direction(scale, u, v);
            for (k, g) in proxies.iter().enumerate() {
                let Some((x, y)) = cameras.project(k, d) else { continue };
                let (x, y) = (x + g.width as f64 / 2.0, y + g.height as f64 / 2.0);
                if x < 0.0 || y < 0.0 || x >= g.width as f64 - 1.0 || y >= g.height as f64 - 1.0 {
                    continue;
                }
                let (x0, y0) = (x as usize, y as usize);
                let (tx, ty) = ((x - x0 as f64) as f32, (y - y0 as f64) as f32);
                let p = |xx: usize, yy: usize| g.data[yy * g.width + xx];
                let val = (p(x0, y0) * (1.0 - tx) + p(x0 + 1, y0) * tx) * (1.0 - ty)
                    + (p(x0, y0 + 1) * (1.0 - tx) + p(x0 + 1, y0 + 1) * tx) * ty;
                sum[oy * out_w + ox] += val;
                count[oy * out_w + ox] += 1;
            }
        }
    }
    let mut ppm = format!("P5\n{out_w} {out_h}\n255\n").into_bytes();
    ppm.extend(sum.iter().zip(&count).map(|(s, c)| {
        if *c == 0 {
            0u8
        } else {
            ((s / f32::from(*c)).clamp(0.0, 1.0) * 255.0) as u8
        }
    }));
    let out = format!("{prefix}-cyl.pgm");
    std::fs::write(&out, ppm).expect("write");
    println!("wrote {out} ({out_w}×{out_h}) in {:?}", t.elapsed());
}