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

//! The face-indexing batch job, off the GUI.
//!
//! Checks every library image for a face-detection run marker, and optionally
//! indexes whatever is missing one.
//!
//!   cargo run -p dr-ui --example face_index -- CATALOG.db THUMBS_DIR [--run DET.onnx EMB.onnx]
//!   cargo run -p dr-ui --example face_index -- CATALOG.db THUMBS_DIR --cluster
//!
//! # Why this exists beside the button in the Identity screen
//!
//! Indexing a real library is hours of work (docs/faces.md §12.2), and the
//! cases where that is worth starting — an overnight pass, a fresh import, a
//! machine left running — are exactly the ones where holding a window open is
//! the wrong shape. The check half is useful on its own: it is cheap, it
//! answers "has face recognition been over all of this", and it distinguishes
//! *not yet indexed* from *waiting on a proxy*, which are different problems
//! with different fixes.
//!
//! The models must have had their input dims frozen first; see
//! `tools/fix-face-model-shapes.sh`.

use std::path::PathBuf;

use dr_catalog::Catalog;
use dr_thumbs::ThumbStore;
use dr_ui::faces::{self, FaceSweepMessage};

const MODEL_ID: &str = "w600k_mbf";

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

    let args: Vec<String> = std::env::args().skip(1).collect();
    if args.len() < 2 {
        eprintln!(
            "usage: face_index CATALOG.db THUMBS_DIR [--run DETECTOR.onnx EMBEDDER.onnx]\n\
             \n\
             With no --run this only reports; nothing is written."
        );
        std::process::exit(2);
    }

    let catalog_path = PathBuf::from(&args[0]);
    let store_dir = PathBuf::from(&args[1]);

    let catalog = match Catalog::open(&catalog_path) {
        Ok(c) => c,
        Err(e) => {
            eprintln!("cannot open catalog {}: {e}", catalog_path.display());
            std::process::exit(1);
        }
    };
    let store = match ThumbStore::open(&store_dir) {
        Ok(s) => s,
        Err(e) => {
            eprintln!("cannot open thumbnail store {}: {e}", store_dir.display());
            std::process::exit(1);
        }
    };

    let audit = match faces::audit(&catalog, &store, MODEL_ID) {
        Ok(a) => a,
        Err(e) => {
            eprintln!("coverage check failed: {e}");
            std::process::exit(1);
        }
    };

    println!("model      {MODEL_ID}");
    println!("images     {}", audit.coverage.images);
    println!(
        "indexed    {} ({:.1}%)",
        audit.coverage.indexed,
        audit.coverage.fraction() * 100.0
    );
    println!("faces      {}", audit.coverage.faces);
    println!(
        "no faces   {}  (indexed, nothing found — the common case)",
        audit.coverage.without_faces
    );
    println!("outstanding {}", audit.coverage.outstanding());
    println!("  ready to index  {}", audit.ready);
    println!("  awaiting proxy  {}", audit.awaiting_proxy);

    // Grouping is a separate step from indexing on purpose: it is a
    // whole-library operation over the embeddings detection produced, and it is
    // worth running *after* a sweep rather than during one (catalog.md §10.2).
    if args.iter().any(|a| a == "--cluster") {
        match dr_ui::faces::recluster(&catalog, MODEL_ID, dr_face::DEFAULT_MERGE_PROBABILITY) {
            Ok((suggested, created)) => {
                println!("\nclustering: {suggested} suggestion(s), {created} new group(s)");
                report_people(&catalog);
            }
            Err(e) => {
                eprintln!("clustering failed: {e}");
                std::process::exit(1);
            }
        }
        return;
    }

    let run = args.iter().position(|a| a == "--run");
    let Some(i) = run else {
        if audit.coverage.is_complete() {
            println!("\nnothing outstanding.");
        } else {
            println!("\npass --run DETECTOR.onnx EMBEDDER.onnx to index the outstanding images.");
        }
        return;
    };

    let (Some(detector), Some(embedder)) = (args.get(i + 1), args.get(i + 2)) else {
        eprintln!("--run needs both a detector and an embedder");
        std::process::exit(2);
    };

    if audit.ready == 0 {
        println!("\nnothing ready to index.");
        if audit.awaiting_proxy > 0 {
            // Worth saying plainly: running this again will not help, because
            // the blocker is in the thumbnail store rather than here.
            println!(
                "{} image(s) are waiting on a proxy — run the thumbnail sweep first.",
                audit.awaiting_proxy
            );
        }
        return;
    }

    println!("\nindexing {} image(s)…", audit.ready);
    let rx = faces::spawn_face_sweep(
        catalog_path,
        store_dir,
        PathBuf::from(detector),
        PathBuf::from(embedder),
        MODEL_ID.to_string(),
        dr_face::DetectOptions::default(),
    );

    let mut seen = 0usize;
    let mut total = 0usize;
    let start = std::time::Instant::now();
    for msg in rx {
        match msg {
            FaceSweepMessage::Total(n) => total = n,
            FaceSweepMessage::Indexed { faces, .. } => {
                seen += 1;
                // One line per image would be thousands of lines; one per
                // twenty-five is enough to see it moving and to estimate.
                if seen.is_multiple_of(25) || faces > 0 {
                    let rate = seen as f64 / start.elapsed().as_secs_f64().max(1e-6);
                    println!(
                        "  {seen}/{total}  {:.2} img/s  ~{:.0} min left",
                        rate,
                        (total.saturating_sub(seen)) as f64 / rate.max(1e-6) / 60.0
                    );
                }
            }
            FaceSweepMessage::Finished {
                images,
                faces,
                failed,
            } => {
                println!(
                    "\ndone: {images} image(s), {faces} face(s), {failed} failed, in {:.0}s",
                    start.elapsed().as_secs_f64()
                );
            }
        }
    }

    if let Ok(a) = faces::audit(&catalog, &store, MODEL_ID) {
        println!("{}", a.summary());
    }
    println!("\nrun again with --cluster to group these faces into people.");
}

/// What the clustering proposed, largest group first.
fn report_people(catalog: &Catalog) {
    let Ok(people) = dr_catalog::faces::people(catalog.connection()) else {
        return;
    };
    if people.is_empty() {
        println!("no groups — too few faces, or none similar enough to group.");
        return;
    }
    println!("\n{} group(s):", people.len());
    for p in people.iter().take(30) {
        let name = if p.name.is_empty() {
            "(unnamed)".to_string()
        } else {
            p.name.clone()
        };
        println!(
            "  {name:<24} {} confirmed, {} suggested",
            p.confirmed_faces, p.suggested_faces
        );
    }
    if people.len() > 30 {
        println!("  … and {} more", people.len() - 30);
    }
}