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

//! Sweep one operation's parameters and write a frame per step.
//!
//! A diagnostic, not part of the product: it exists to show what a control
//! actually does to a photograph, one parameter at a time, so a new node can
//! be looked at rather than reasoned about.
//!
//! ```sh
//! cargo run -p dr-gpu --release --example sweep -- IMG.CR2 out/ colour_grading 21
//! ```
//!
//! It names no operation. The op id arrives as a string, the parameters and
//! their ranges come from the graph's own capabilities, and a node declared
//! yesterday sweeps on the same terms as one that shipped a year ago — which
//! is the property `ops/README.md` promises and the reason this is one example
//! rather than one per node.
//!
//! PPM out, like `develop.rs`, so it needs no encoder dependency; the caller
//! turns them into whatever it wants.

use dr_gpu::{AdjustPass, DemosaicedImage, Demosaicer, GpuContext};
use dr_pipeline::{Affects, EditGraph, OpId, ParamId, ParamKind};
use dr_types::ColourSpace;

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

    let mut args = std::env::args().skip(1);
    let (Some(input), Some(out_dir), Some(op)) = (args.next(), args.next(), args.next()) else {
        eprintln!("usage: sweep <file.cr2> <out_dir> <op_id> [steps]");
        eprintln!("       sweep <file.cr2> <out_dir> --list");
        std::process::exit(2);
    };
    let steps: usize = args
        .next()
        .and_then(|s| s.parse().ok())
        .filter(|n| *n >= 2)
        .unwrap_or(21);
    // Sweep one named parameter rather than all of them.
    let only: Option<String> = args.next();

    let ctx = pollster::block_on(GpuContext::new_headless()).expect("gpu");
    let image = load(&ctx, &input);

    let mut graph = EditGraph::default_chain();

    // `--list` prints every op and parameter with its range, which is how the
    // caller learns what there is to sweep without this file holding a list
    // that would go stale.
    if op == "--list" {
        for cap in graph.capabilities() {
            println!("{}", cap.id.0);
            for p in &cap.params {
                match p.kind {
                    ParamKind::Scalar { min, max, .. } => {
                        println!("    {:<20} {min} .. {max}   default {}", p.id.0, p.default)
                    }
                    ParamKind::Bool => println!("    {:<20} bool", p.id.0),
                    ParamKind::Enum { ref variants } => {
                        println!("    {:<20} enum, {} variants", p.id.0, variants.len())
                    }
                }
            }
        }
        return;
    }

    std::fs::create_dir_all(&out_dir).expect("create out dir");

    // `OpId`/`ParamId` hold `&'static str`, and an argument is not static.
    // Leaking is right rather than expedient here: the ids live as long as the
    // graph does, and this process exits immediately after.
    let op_id = OpId(Box::leak(op.clone().into_boxed_str()));

    let cap = graph
        .capabilities()
        .into_iter()
        .find(|c| c.id == op_id)
        .unwrap_or_else(|| {
            eprintln!("no operation `{op}` — try --list");
            std::process::exit(1);
        });

    // Bounded output rather than full sensor resolution. This is the proxy
    // path FR-DSP-1 already renders through, so it is the same code the
    // develop view uses — and a 25 MP frame would be a 75 MB PPM, times
    // several hundred frames in one sweep.
    let (full_w, full_h) = image.size();
    let longest = std::env::var("SWEEP_MAX_PX")
        .ok()
        .and_then(|s| s.parse::<u32>().ok())
        .unwrap_or(1100);
    let scale = (longest as f32 / full_w.max(full_h) as f32).min(1.0);
    let w = ((full_w as f32 * scale) as u32).max(1);
    let h = ((full_h as f32 * scale) as u32).max(1);
    println!("rendering  {w} x {h} (from {full_w} x {full_h})");

    let mut adjust = AdjustPass::new(&ctx);

    // The reference frame: every parameter at its default. Written once so a
    // viewer can see what the sweep is departing from.
    render_to(&mut adjust, &image, &graph, w, h, &out_dir, "neutral", 0, 0.0);

    // Parameters held away from their default for the duration, as
    // `SWEEP_HOLD=shadow_strength=70,midtone_hue=210`.
    //
    // Needed because a parameter is not always meaningful alone. Where a
    // `presentation:` block groups several into one conceptual control — a
    // hue and the strength behind it — sweeping one with the other at its
    // default renders the same frame every time, which looks like a broken
    // node rather than a correctly declared neutral.
    let hold = std::env::var("SWEEP_HOLD").unwrap_or_default();
    for clause in hold.split(',').filter(|c| !c.trim().is_empty()) {
        let Some((name, value)) = clause.split_once('=') else {
            eprintln!("SWEEP_HOLD wants name=value, got `{clause}`");
            std::process::exit(2);
        };
        let value: f32 = value.trim().parse().expect("hold value");
        let name = Box::leak(name.trim().to_string().into_boxed_str());
        graph.set_param(op_id, ParamId(name), value);
        println!("holding    {name} = {value}");
    }

    for p in &cap.params {
        if only.as_deref().is_some_and(|o| o != p.id.0) {
            continue;
        }
        let ParamKind::Scalar { min, max, .. } = p.kind else {
            eprintln!("skipping {} — only scalars sweep meaningfully", p.id.0);
            continue;
        };
        let param_id = ParamId(Box::leak(p.id.0.to_string().into_boxed_str()));

        for i in 0..steps {
            let t = i as f32 / (steps - 1) as f32;
            let value = min + (max - min) * t;
            graph.set_param(op_id, param_id, value);
            render_to(
                &mut adjust,
                &image,
                &graph,
                w,
                h,
                &out_dir,
                p.id.0,
                i,
                value,
            );
        }
        // Back to default before the next parameter, so each sweep is of one
        // control rather than of everything tried so far.
        graph.set_param(op_id, param_id, p.default);
    }
}

#[allow(clippy::too_many_arguments)]
fn render_to(
    adjust: &mut AdjustPass,
    image: &dr_gpu::DemosaicedImage,
    graph: &EditGraph,
    w: u32,
    h: u32,
    dir: &str,
    name: &str,
    index: usize,
    value: f32,
) {
    // The detail path, always. A neighbourhood node — dehaze, clarity,
    // sharpening — runs as its own dispatch after the fused pass, and the
    // fused path refuses a shader composed with one rather than rendering it
    // wrongly. `render_detailed` falls through to the plain path when the
    // chain has no detail stage, so this one call serves both kinds of node
    // and the example never has to know which it was handed.
    let shader = graph.compose_for(ColourSpace::Srgb);
    let scale = graph.render_scale(image.size(), (w, h));
    let detail = graph.compose_detail_for(scale.full_size(), scale.render_size(), ColourSpace::Srgb);
    let key = graph.invalidation().through(Affects::Colour);
    adjust
        .render_detailed(image, &shader, w, h, None, &detail, key)
        .expect("adjust");
    let (pixels, pw, ph) = adjust.export_pixels().expect("readback");

    let path = format!("{dir}/{name}_{index:03}.ppm");
    write_ppm(&path, &pixels, pw, ph);

    // The value goes beside the frame rather than into the filename: a caption
    // wants "-37.5", and a filename that carried it would need escaping and
    // would sort wrongly.
    let meta = format!("{dir}/{name}_{index:03}.txt");
    std::fs::write(meta, format!("{name} {value:.4}\n")).expect("write value");
    println!("{name}[{index}] = {value:.4}");
}

fn write_ppm(path: &str, rgba: &[u8], w: u32, h: u32) {
    use std::io::Write;
    let mut out = Vec::with_capacity((w * h * 3) as usize + 32);
    out.extend_from_slice(format!("P6\n{w} {h}\n255\n").as_bytes());
    for px in rgba.chunks_exact(4) {
        out.extend_from_slice(&px[..3]);
    }
    std::fs::File::create(path)
        .expect("create output")
        .write_all(&out)
        .expect("write output");
}

/// Load either a RAW or an already-rendered image.
///
/// A JPEG takes the path `DemosaicedImage::from_rgba8` documents: no CFA to
/// interpolate, identity colour matrix, neutral white balance, and the shader
/// linearises the gamma-encoded pixels. The controls all still work; their
/// neutral is "as the camera left it" rather than "as the sensor recorded it",
/// which is worth knowing when reading a sweep made from one.
fn load(ctx: &GpuContext, path: &str) -> DemosaicedImage {
    let bytes = std::fs::read(path).expect("read file");
    match dr_decode::probe(&bytes) {
        Some(dr_types::Format::Jpeg) => {
            let p = dr_decode::decode_jpeg(&bytes).expect("decode jpeg");
            println!("loaded     {} x {} (rendered, not raw)", p.width, p.height);
            DemosaicedImage::from_rgba8(ctx, &p.rgba, p.width, p.height).expect("upload")
        }
        _ => {
            let raw = dr_decode::decode(&bytes).expect("decode raw");
            println!("loaded     {} x {} (raw)", raw.crop.width, raw.crop.height);
            let demosaic = Demosaicer::new(ctx).expect("demosaicer");
            demosaic.run(&raw).expect("demosaic")
        }
    }
}