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

//! Render a RAW file through the full pipeline and write a PPM.
//!
//! The end-to-end check: decode → demosaic → adjust → display encode, on a
//! real file rather than a synthetic fixture. Unit tests prove each stage in
//! isolation; this proves they compose into an image a person would accept.
//!
//! ```sh
//! cargo run -p dr-gpu --example develop --features readback -- IMG.CR2 out.ppm
//! ```
//!
//! PPM because it needs no encoder dependency and every image viewer reads
//! it. This is a diagnostic, not the export path (FR-EXP-*).

use dr_gpu::{AdjustPass, Demosaicer, GpuContext};
use dr_pipeline::ops::{
    blacks_whites, brilliance, colour_mixer, contrast, curve, exposure, highlights_shadows,
    vibrance, white_balance,
};
use dr_pipeline::{EditGraph, ParamId};

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

    let mut args = std::env::args().skip(1);
    let Some(input) = args.next() else {
        eprintln!("usage: develop <file.cr2> [out.ppm] [preset]");
        eprintln!("  preset: neutral (default) | punchy | recover");
        std::process::exit(2);
    };
    let output = args.next().unwrap_or_else(|| "develop.ppm".into());
    let preset = args.next().unwrap_or_else(|| "neutral".into());

    let bytes = std::fs::read(&input).expect("read file");

    let t0 = std::time::Instant::now();
    let raw = dr_decode::decode(&bytes).expect("decode");
    let decode_ms = t0.elapsed().as_secs_f32() * 1000.0;
    println!(
        "decoded    {} × {} ({:?}), {decode_ms:.0} ms",
        raw.crop.width, raw.crop.height, raw.cfa_pattern
    );

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

    let t1 = std::time::Instant::now();
    let demosaicer = Demosaicer::new(&ctx).expect("demosaicer");
    let image = demosaicer.run(&raw).expect("demosaic");
    ctx.device
        .poll(wgpu::PollType::wait_indefinitely())
        .expect("poll");
    println!("demosaiced {:.0} ms", t1.elapsed().as_secs_f32() * 1000.0);

    // Build an edit. The presets exist so the output can be eyeballed for
    // each operation actually doing something, not merely compiling.
    let mut graph = EditGraph::default_chain();
    match preset.as_str() {
        "punchy" => {
            graph.set_param(exposure::ID, exposure::EXPOSURE, 0.3);
            graph.set_param(
                highlights_shadows::ID,
                highlights_shadows::HIGHLIGHTS,
                -40.0,
            );
            graph.set_param(highlights_shadows::ID, highlights_shadows::SHADOWS, 30.0);
            graph.set_param(blacks_whites::ID, blacks_whites::BLACKS, -20.0);
            graph.set_param(blacks_whites::ID, blacks_whites::WHITES, 25.0);
            graph.set_param(vibrance::ID, vibrance::VIBRANCE, 35.0);
        }
        "recover" => {
            graph.set_param(exposure::ID, exposure::EXPOSURE, -0.5);
            graph.set_param(
                highlights_shadows::ID,
                highlights_shadows::HIGHLIGHTS,
                -80.0,
            );
            graph.set_param(highlights_shadows::ID, highlights_shadows::SHADOWS, 60.0);
            graph.set_param(brilliance::ID, brilliance::BRILLIANCE, 40.0);
            graph.set_param(white_balance::ID, white_balance::TEMPERATURE, 15.0);
        }
        // Contrast alone, so its effect can be judged without anything else
        // moving.
        "contrast" => {
            graph.set_param(contrast::ID, contrast::CONTRAST, 60.0);
        }
        "flat" => {
            graph.set_param(contrast::ID, contrast::CONTRAST, -60.0);
        }
        // The mixer, pushed hard on the two things this scene actually has:
        // green vegetation and grey-blue rock.
        "mixer" => {
            graph.set_param(colour_mixer::ID, ParamId("green_sat"), 80.0);
            graph.set_param(colour_mixer::ID, ParamId("green_hue"), -40.0);
            graph.set_param(colour_mixer::ID, ParamId("chartreuse_sat"), 60.0);
            graph.set_param(colour_mixer::ID, ParamId("azure_lum"), -50.0);
        }
        // One band only, to check the weighting really is selective rather
        // than affecting the whole image.
        "mixer_one" => {
            graph.set_param(colour_mixer::ID, ParamId("green_sat"), 100.0);
        }
        // A classic S-curve: shadows down, highlights up, mid held.
        "curve_s" => {
            graph.set_param(curve::ID, curve::P1_Y, 0.15);
            graph.set_param(curve::ID, curve::P3_Y, 0.85);
        }
        // The inverse, a film-like lifted-shadow look.
        "curve_lift" => {
            graph.set_param(curve::ID, curve::P0_Y, 0.12);
            graph.set_param(curve::ID, curve::P1_Y, 0.32);
        }
        _ => {}
    }

    let shader = graph.compose();
    println!(
        "shader     {} active op(s), {} uniform floats, structure {:016x}",
        shader.source.matches("---- ").count(),
        shader.uniforms.len(),
        shader.structure_hash
    );

    let mut adjust = AdjustPass::new(&ctx);
    let (w, h) = image.size();

    let t2 = std::time::Instant::now();
    adjust.render(&image, &shader, w, h).expect("adjust");
    ctx.device
        .poll(wgpu::PollType::wait_indefinitely())
        .expect("poll");
    println!("adjusted   {:.2} ms", t2.elapsed().as_secs_f32() * 1000.0);

    // Time a second render with only a value changed: this is the slider
    // path, and it must not recompile.
    graph.set_param(exposure::ID, exposure::EXPOSURE, 0.31);
    let again = graph.compose();
    let t3 = std::time::Instant::now();
    adjust.render(&image, &again, w, h).expect("adjust");
    ctx.device
        .poll(wgpu::PollType::wait_indefinitely())
        .expect("poll");
    println!(
        "re-render  {:.2} ms ({} pipeline(s) compiled)",
        t3.elapsed().as_secs_f32() * 1000.0,
        adjust.cached_pipelines()
    );

    let (pixels, pw, ph) = adjust.read_output().expect("readback");

    // Sanity: an all-black or all-white result means something upstream
    // failed silently, and it is far easier to see here than in a viewer.
    let mut sum = 0u64;
    let mut min = 255u8;
    let mut max = 0u8;
    for px in pixels.chunks_exact(4) {
        let l = px[0].max(px[1]).max(px[2]);
        sum += u64::from(l);
        min = min.min(l);
        max = max.max(l);
    }
    let mean = sum as f64 / (pixels.len() / 4) as f64;
    println!("levels     min {min}, mean {mean:.1}, max {max}");
    if max == 0 {
        eprintln!("WARNING: the image is entirely black");
    }

    write_ppm(&output, &pixels, pw, ph);
    println!("wrote      {output} ({pw} × {ph})");
}

/// Write binary PPM (P6): a three-line header then RGB triples.
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");
}