DarkRoom/core/dr-export/examples/export.rs

//! Export a real file, end to end, from a real image.
//!
//!   cargo run -p dr-export --example export -- <file.jpg|file.cr2> [out-dir]
//!
//! Deliberately the *whole* path and not a unit test of the encoder: decode,
//! demosaic or upload, run the develop chain on the GPU at full resolution,
//! read the result back through `AdjustPass::export_pixels`, resize, sharpen,
//! encode, and write. A test can prove the JPEG has the right magic bytes; it
//! cannot tell anyone whether the picture came out looking like the picture.

use std::path::PathBuf;

use dr_export::{export, Frame, NameContext};
use dr_gpu::{AdjustPass, DemosaicedImage, Demosaicer, GpuContext};
use dr_pipeline::EditGraph;
use dr_types::{ColourSpace, ExportFormat, ExportSettings, OutputSharpening, SizingMode};

fn main() {
    env_logger::Builder::from_env(env_logger::Env::default().default_filter_or("info,wgpu=warn"))
        .init();

    let mut args = std::env::args().skip(1);
    let Some(input) = args.next() else {
        eprintln!("usage: export <file.jpg|file.cr2> [out-dir]");
        std::process::exit(2);
    };
    let out_dir = PathBuf::from(args.next().unwrap_or_else(|| ".".into()));
    let input = PathBuf::from(input);

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

    // Decode. A RAW goes through the demosaicer; a JPEG is already RGB and
    // takes the same path every operation after the sensor stage does.
    let bytes = std::fs::read(&input).expect("read input");
    // From content, not from the extension — dr-decode is emphatic that an
    // extension is only a hint. Its own `probe` reports a crate-private
    // `Format`, so the SOI marker is checked directly here rather than
    // widening that API for an example.
    let is_jpeg = bytes.starts_with(&[0xFF, 0xD8, 0xFF]);
    let source = if !is_jpeg {
        let raw = dr_decode::decode(&bytes).expect("decode raw");
        let demosaicer = Demosaicer::new(&ctx).expect("demosaicer");
        demosaicer.run(&raw).expect("demosaic")
    } else {
        let (rgba, w, h) = decode_jpeg(&bytes);
        DemosaicedImage::from_rgba8(&ctx, &rgba, w, h).expect("upload")
    };

    // An edit worth seeing in the output, so a broken pipeline is obvious
    // rather than subtle.
    let mut graph = EditGraph::default_chain();
    graph.set_param(
        dr_pipeline::ops::exposure::ID,
        dr_pipeline::ops::exposure::EXPOSURE,
        0.35,
    );
    graph.set_param(
        dr_pipeline::ops::contrast::ID,
        dr_pipeline::ops::contrast::CONTRAST,
        18.0,
    );
    graph.set_param(
        dr_pipeline::ops::saturation::ID,
        dr_pipeline::ops::saturation::SATURATION,
        12.0,
    );

    // Full resolution, not the viewport (FR-EXP-9). This is the one thing an
    // export must not economise on.
    let (sw, sh) = source.size();
    let (fw, fh) = graph.output_size(sw, sh);
    println!("source {sw}×{sh}, framed {fw}×{fh}");

    // The output space is chosen *here*, before the render, because that is
    // where it takes effect: the primaries conversion and the encode are the
    // last two lines of the generated shader (FR-EXP-2). Asking for it at the
    // encoder would be too late — the pixels would already be clipped.
    let space = ColourSpace::DisplayP3;

    let mut adjust = AdjustPass::new(&ctx);
    let shader = graph.compose_for(space);
    let t = std::time::Instant::now();
    adjust.render(&source, &shader, fw, fh).expect("render");
    let (pixels, w, h) = adjust.export_pixels().expect("read back");
    println!(
        "rendered {w}×{h} in {:.0} ms as {}",
        t.elapsed().as_secs_f32() * 1000.0,
        space.label()
    );

    let frame = Frame::in_space(w, h, pixels, space).expect("well-formed frame");

    let stem = input
        .file_stem()
        .map(|s| s.to_string_lossy().into_owned())
        .unwrap_or_else(|| "export".into());

    // One of each format, so the run exercises every encoder that exists.
    for (format, sizing, sharpening) in [
        (
            ExportFormat::Jpeg,
            SizingMode::Original,
            OutputSharpening::None,
        ),
        (
            ExportFormat::Jpeg,
            SizingMode::LongEdge(1200),
            OutputSharpening::Screen,
        ),
        (
            ExportFormat::Png,
            SizingMode::LongEdge(600),
            OutputSharpening::Screen,
        ),
        (
            ExportFormat::Tiff8,
            SizingMode::Percentage(25),
            OutputSharpening::MattePaper,
        ),
        (
            ExportFormat::Tiff16,
            SizingMode::Percentage(25),
            OutputSharpening::MattePaper,
        ),
    ] {
        let settings = ExportSettings {
            format,
            sizing,
            sharpening,
            colour_space: space,
            filename_template: "{name}-{dimensions}".into(),
            ..Default::default()
        };

        // The size has to be known before the name, because `{dimensions}` is
        // part of it — which is why sizing is resolved here and not inside
        // `export`.
        let (tw, th) = dr_export::target_size(w, h, sizing, settings.allow_upscaling);
        let ctx = NameContext {
            source_stem: &stem,
            sequence: 1,
            date: "",
            width: tw,
            height: th,
            preset: "",
        };
        let name = dr_export::resolve_name(
            &settings.filename_template,
            &ctx,
            format,
            settings.collision,
            &|n| out_dir.join(n).exists(),
        )
        .expect("a free name");

        let t = std::time::Instant::now();
        let out = export(&frame, &settings, name).expect("export");
        let path = out_dir.join(&out.name);
        std::fs::write(&path, &out.bytes).expect("write");
        println!(
            "{:>10} {:>5}×{:<5} {:>8} KB  {:>5.0} ms  {}",
            format.label(),
            out.width,
            out.height,
            out.bytes.len() / 1024,
            t.elapsed().as_secs_f32() * 1000.0,
            path.display()
        );
    }
}

fn decode_jpeg(bytes: &[u8]) -> (Vec<u8>, u32, u32) {
    let mut decoder = zune_jpeg::JpegDecoder::new(bytes);
    let pixels = decoder.decode().expect("decode jpeg");
    let info = decoder.info().expect("jpeg info");
    let (w, h) = (u32::from(info.width), u32::from(info.height));

    // zune gives RGB; the GPU upload wants RGBA.
    let mut rgba = Vec::with_capacity((w * h * 4) as usize);
    for px in pixels.chunks_exact(3) {
        rgba.extend_from_slice(&[px[0], px[1], px[2], 255]);
    }
    (rgba, w, h)
}