DarkRoom/core/dr-gpu/examples/segment.rs

//! Segment an image and write the granularity ladder as false-coloured PPMs.
//!
//! The whole point of S15 step 2 (docs/dev/segmentation.md §11): look at the
//! ladder and decide whether clicking through it would land on the things a
//! person means. No amount of design settles that — the pictures do.
//!
//! ```sh
//! cargo run -p dr-gpu --example segment --features readback -- IMG.CR2
//! cargo run -p dr-gpu --example segment --features readback -- synthetic
//! ```
//!
//! PPM for the same reason `develop` uses it: no encoder dependency, and
//! every viewer reads it. This is a diagnostic, not an export path.

use dr_gpu::{DemosaicedImage, Demosaicer, GpuContext, SegmentOptions, SegmentPass};
use dr_segment::{MergeTree, RegionField};

/// The ladder the example dumps. Chosen to span "far too fine to be useful"
/// through "one or two objects", because both ends are informative: if no rung
/// looks right, the gradient is wrong rather than the ladder being too coarse.
const LEVELS: [usize; 6] = [2000, 800, 300, 120, 50, 16];

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

    let mut args = std::env::args().skip(1);
    let Some(input) = args.next() else {
        eprintln!("usage: segment <file.raw|synthetic> [out-prefix] [blur-radius]");
        std::process::exit(2);
    };
    let prefix = args.next().unwrap_or_else(|| "segment".into());
    let blur_radius = args
        .next()
        .and_then(|s| s.parse().ok())
        .unwrap_or(SegmentOptions::default().blur_radius);

    let ctx = pollster::block_on(GpuContext::new_headless()).expect("gpu context");
    println!("gpu        {}", ctx.adapter_name());

    let source = if input == "synthetic" {
        let (w, h) = (1200, 800);
        println!("source     synthetic {w} × {h}");
        DemosaicedImage::from_rgba8(&ctx, &synthetic(w, h), w, h).expect("synthetic source")
    } else {
        let bytes = std::fs::read(&input).expect("read file");
        let raw = dr_decode::decode(&bytes).expect("decode");
        println!("source     {} × {}", raw.crop.width, raw.crop.height);
        Demosaicer::new(&ctx)
            .expect("demosaicer")
            .run(&raw)
            .expect("demosaic")
    };

    let opts = SegmentOptions {
        blur_radius,
        ..Default::default()
    };
    let pass = SegmentPass::new(&ctx).expect("segment pass");

    let t0 = std::time::Instant::now();
    let seg = pass.run(&source, opts).expect("segment");
    let (w, h) = seg.size();
    let field = seg.read_field().expect("read field");
    let gpu_ms = t0.elapsed().as_secs_f32() * 1000.0;

    let t1 = std::time::Instant::now();
    let tree = MergeTree::build(&field);
    let tree_ms = t1.elapsed().as_secs_f32() * 1000.0;

    // M6, roughly: the readback is in `gpu_ms` and would not be there in a
    // shipping build, so this over-reports the GPU half rather than under.
    println!("proxy      {w} × {h}, blur radius {blur_radius}");
    println!("basins     {}", field.region_count);
    println!("boundaries {}", field.adjacency.len());
    println!("merges     {}", tree.merges.len());
    println!("segment    {gpu_ms:.0} ms (includes readback)");
    println!("hierarchy  {tree_ms:.1} ms");

    if let (Some(first), Some(last)) = (tree.merges.first(), tree.merges.last()) {
        println!("saddles    {:.4} … {:.4}", first.saddle, last.saddle);
    }

    for level in LEVELS {
        if level > field.region_count {
            println!("skip       {level} (only {} basins)", field.region_count);
            continue;
        }
        let grouping = tree.cut_to(level);
        let pixels = field.apply(&grouping);
        let groups = grouping.iter().max().map(|m| m + 1).unwrap_or(0);
        let path = format!("{prefix}-{level:04}.ppm");
        write_ppm(&path, &false_colour(&pixels, &field), w, h);
        println!("wrote      {path} ({groups} regions)");
    }

    // The boundaries alone, which is what a snapped contour would cling to.
    let path = format!("{prefix}-edges.ppm");
    write_ppm(&path, &boundaries(&field), w, h);
    println!("wrote      {path}");
}

/// A distinct colour per region.
///
/// Hashed from the id rather than sampled from the image: two adjacent
/// regions that happen to look alike are exactly the case worth seeing, and
/// mean colours would hide it.
fn false_colour(pixels: &[u32], field: &RegionField) -> Vec<u8> {
    let _ = field;
    let mut out = Vec::with_capacity(pixels.len() * 3);
    for &g in pixels {
        // Cheap integer hash — golden-ratio multiply, then spread the bits
        // across three channels.
        let mut x = g.wrapping_mul(2_654_435_761);
        x ^= x >> 15;
        out.push((x & 0xff) as u8);
        out.push(((x >> 8) & 0xff) as u8);
        out.push(((x >> 16) & 0xff) as u8);
    }
    out
}

/// White where two regions meet, black elsewhere.
fn boundaries(field: &RegionField) -> Vec<u8> {
    let (w, h) = (field.width, field.height);
    let mut out = vec![0u8; w * h * 3];
    for y in 0..h {
        for x in 0..w {
            let i = y * w + x;
            let edge = (x + 1 < w && field.labels[i] != field.labels[i + 1])
                || (y + 1 < h && field.labels[i] != field.labels[i + w]);
            if edge {
                out[i * 3] = 255;
                out[i * 3 + 1] = 255;
                out[i * 3 + 2] = 255;
            }
        }
    }
    out
}

fn write_ppm(path: &str, rgb: &[u8], width: u32, height: u32) {
    use std::io::Write as _;
    let mut f = std::io::BufWriter::new(std::fs::File::create(path).expect("create ppm"));
    write!(f, "P6\n{width} {height}\n255\n").expect("ppm header");
    f.write_all(rgb).expect("ppm body");
}

/// A test image with the failure modes the corpus is meant to provoke, so the
/// example is runnable before anyone has traced a single ground-truth mask.
///
/// Deliberately includes a soft gradient boundary and a noisy patch: those are
/// where a watershed either earns its place or shatters, and a synthetic image
/// of clean shapes would flatter it.
fn synthetic(w: u32, h: u32) -> Vec<u8> {
    let mut px = Vec::with_capacity((w * h * 4) as usize);
    for y in 0..h {
        for x in 0..w {
            let fx = x as f32 / w as f32;
            let fy = y as f32 / h as f32;

            // A smooth vertical gradient — the low-contrast boundary case.
            let mut r = 40.0 + 120.0 * fy;
            let mut g = 60.0 + 100.0 * fy;
            let mut b = 110.0 + 90.0 * fy;

            // A hard-edged disc: the control case.
            let d = ((fx - 0.3).powi(2) + (fy - 0.45).powi(2)).sqrt();
            if d < 0.16 {
                r = 210.0;
                g = 90.0;
                b = 60.0;
            }

            // A soft-edged disc: where the ladder should merge late.
            let d2 = ((fx - 0.68).powi(2) + (fy - 0.6).powi(2)).sqrt();
            let t = (1.0 - (d2 / 0.18)).clamp(0.0, 1.0);
            r = r * (1.0 - t) + 90.0 * t;
            g = g * (1.0 - t) + 170.0 * t;
            b = b * (1.0 - t) + 110.0 * t;

            // A noisy corner: the case pre-smoothing exists for.
            if fx > 0.82 && fy < 0.22 {
                let n = ((x * 7919 + y * 104_729) % 97) as f32 / 97.0;
                r += (n - 0.5) * 90.0;
                g += (n - 0.5) * 90.0;
                b += (n - 0.5) * 90.0;
            }

            px.push(r.clamp(0.0, 255.0) as u8);
            px.push(g.clamp(0.0, 255.0) as u8);
            px.push(b.clamp(0.0, 255.0) as u8);
            px.push(255);
        }
    }
    px
}