DarkRoom/ui/dr-ui/examples/face_native.rs

//! TRACES: FR-CULL-8 | FR-EXP-9
//! Measure what indexing at native resolution is actually worth.
//!
//!   cargo run -p dr-ui --example face_native -- DET.onnx EMB.onnx [--export DIR] FILE [FILE…]
//!
//! Renders each file once at native resolution, then indexes it twice: the way
//! FR-CULL-8 now specifies, and the way it used to be done — everything, both
//! stages, from a 1024px proxy. Prints the faces found and the `crop_px` each
//! run gave the embedder.
//!
//! It exists because the case for the change was made from `crop_px` readings
//! taken out of a catalog after the fact. That is evidence about what happened;
//! this is evidence about what the new code does, on the same photographs, with
//! nothing between the two runs but the resolution.
//!
//! `--export DIR` also writes each native render out as a JPEG, so the
//! model-free tooling in `dr-face`'s examples — `eyes` above all — can be
//! run over native pixels rather than proxies (docs/faces.md §17.4).
//!
//! The models must have had their input dims frozen first; see
//! `tools/fix-face-model-shapes.sh`.

use std::path::PathBuf;

const MODEL_ID: &str = "w600k_mbf";

fn main() {
    env_logger::init();

    let mut args: Vec<String> = std::env::args().skip(1).collect();
    let export = args.iter().position(|a| a == "--export").map(|i| {
        args.remove(i);
        PathBuf::from(args.remove(i))
    });
    if args.len() < 3 {
        eprintln!("usage: face_native DETECTOR.onnx EMBEDDER.onnx [--export DIR] FILE [FILE…]");
        std::process::exit(2);
    }
    if let Some(d) = &export {
        std::fs::create_dir_all(d).expect("export dir");
    }
    let (detector_model, embedder_model) = (PathBuf::from(&args[0]), PathBuf::from(&args[1]));

    let Some(gpu) = pollster::block_on(dr_gpu::GpuContext::new_headless()).ok() else {
        eprintln!("no GPU adapter; a native render needs one");
        std::process::exit(1);
    };

    let mut detector = match dr_face::Detector::from_path(&detector_model) {
        Ok(d) => d,
        Err(e) => {
            eprintln!("detector: {e}");
            std::process::exit(1);
        }
    };
    let mut embedder =
        match dr_face::Embedder::from_path(&embedder_model, dr_face::ModelId::new(MODEL_ID)) {
            Ok(e) => e,
            Err(e) => {
                eprintln!("embedder: {e}");
                std::process::exit(1);
            }
        };
    let options = dr_face::DetectOptions::default();

    println!(
        "{:<28} {:>11} {:>17} {:>17}",
        "file", "native", "native faces/px", "1024 faces/px"
    );

    let (mut n_native, mut n_proxy) = (0usize, 0usize);
    let (mut px_native, mut px_proxy) = (0.0f32, 0.0f32);

    for path in &args[2..] {
        let bytes = match std::fs::read(path) {
            Ok(b) => b,
            Err(e) => {
                println!("{path}: cannot read: {e}");
                continue;
            }
        };
        let frame = match dr_ui::render_native(&gpu, &bytes) {
            Ok(f) => f,
            Err(e) => {
                println!("{path}: cannot render: {e}");
                continue;
            }
        };
        let (w, h) = (frame.width as usize, frame.height as usize);
        if let Some(dir) = &export {
            let stem = std::path::Path::new(path)
                .file_stem()
                .map(|s| s.to_string_lossy().into_owned())
                .unwrap_or_default();
            match dr_thumbs::codec::encode_rgba(frame.width, frame.height, &frame.rgba) {
                Ok(bytes) => std::fs::write(dir.join(format!("{stem}.jpg")), bytes).expect("write"),
                Err(e) => println!("{path}: export: {e}"),
            }
        }

        let native = dr_ui::faces::index_native(
            &mut detector,
            &mut embedder,
            None,
            &frame.rgba,
            w,
            h,
            &options,
        )
        .unwrap_or_default();

        // The old path, reproduced exactly: one buffer at 1024, used for both
        // detection and the crop.
        let mut small = dr_decode::Preview {
            width: frame.width,
            height: frame.height,
            rgba: frame.rgba.clone(),
        };
        small.downscale_to(dr_thumbs::ThumbSize::Large.edge());
        let proxy =
            dr_ui::faces::index_preview(&mut detector, &mut embedder, None, &small, &options)
                .map(|(f, _)| f)
                .unwrap_or_default();

        let mean = |v: &[dr_catalog::faces::DetectedFace]| {
            if v.is_empty() {
                0.0
            } else {
                v.iter().map(|f| f.crop_px).sum::<f32>() / v.len() as f32
            }
        };
        let name = std::path::Path::new(path)
            .file_name()
            .map(|s| s.to_string_lossy().into_owned())
            .unwrap_or_else(|| path.clone());

        println!(
            "{:<28} {:>11} {:>10} /{:>5.0} {:>10} /{:>5.0}",
            name,
            format!("{}x{}", frame.width, frame.height),
            native.len(),
            mean(&native),
            proxy.len(),
            mean(&proxy),
        );

        n_native += native.len();
        n_proxy += proxy.len();
        px_native += native.iter().map(|f| f.crop_px).sum::<f32>();
        px_proxy += proxy.iter().map(|f| f.crop_px).sum::<f32>();
    }

    println!(
        "\ntotal: native {n_native} face(s), mean crop {:.0}px | \
         1024 proxy {n_proxy} face(s), mean crop {:.0}px",
        if n_native == 0 {
            0.0
        } else {
            px_native / n_native as f32
        },
        if n_proxy == 0 {
            0.0
        } else {
            px_proxy / n_proxy as f32
        },
    );
}