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

//! List each RAW file's hot and dead photosite candidates (docs/dev/sensor-health.md).
//!
//! Reads paths on stdin and prints one JSON line per file: its capture
//! conditions and every photosite [`Demosaicer::find_hot_pixels`] flags, as
//! `[x, y, value, hot]` in sensor coordinates. Which candidates are defects
//! is a question across frames, so this answers nothing on its own.
//!
//! ```sh
//! find ~/Pictures -name '*.CR2' | cargo run --release -p dr-gpu --example sensor_scan
//! ```

use std::io::{BufRead, Write};

use dr_gpu::{Demosaicer, GpuContext};

fn main() {
    if let Some(list) = std::env::args().skip_while(|a| a != "--probe").nth(1) {
        return probe(&list);
    }
    let ctx = pollster::block_on(GpuContext::new_headless()).expect("a GPU for the hot-pixel pass");
    let demosaicer = Demosaicer::new(&ctx).expect("demosaicer");
    for line in std::io::stdin().lock().lines() {
        let path = line.expect("stdin");
        match scan(&path, &demosaicer) {
            Ok(json) => println!("{json}"),
            Err(e) => eprintln!("fail\t{path}\t{e}"),
        }
        std::io::stdout().flush().ok();
    }
}

fn scan(path: &str, demosaicer: &Demosaicer) -> Result<String, String> {
    let bytes = std::fs::read(path).map_err(|e| e.to_string())?;
    let raw = dr_decode::decode(&bytes).map_err(|e| e.to_string())?;
    let meta = dr_decode::metadata(&bytes).map_err(|e| e.to_string())?;
    let sites = demosaicer
        .find_hot_pixels(&raw)
        .map_err(|e| e.to_string())?;
    let opt = |v: Option<f64>| v.map_or("null".to_string(), |v| v.to_string());
    let list: Vec<String> = sites
        .iter()
        .map(|s| {
            let v = raw.data[(s.y * raw.width + s.x) as usize];
            format!("[{},{},{},{}]", s.x, s.y, v, u8::from(s.hot))
        })
        .collect();
    Ok(format!(
        "{{\"path\":{:?},\"model\":{:?},\"captured\":{},\"iso\":{},\"shutter\":{},\"white\":{},\"black\":{:?},\"crop\":[{},{},{},{}],\"sites\":[{}]}}",
        path,
        format!("{} {}", raw.make, raw.model),
        meta.captured_at.map_or("null".to_string(), |t| t.to_string()),
        opt(meta.iso.map(f64::from)),
        opt(meta.shutter.map(f64::from)),
        raw.white_level,
        raw.black_level,
        raw.crop.x,
        raw.crop.y,
        raw.crop.width,
        raw.crop.height,
        list.join(","),
    ))
}

/// `--probe COORDS`: for each path on stdin, each `x y` line of COORDS as
/// `[value, same-colour neighbour max, median]` over black, on the CPU. A
/// probe asks whether a photosite stood out in a frame where it would have
/// been visible, which the scan's verdict cannot say: a frame that does not
/// flag a defect may only have been too bright around it.
fn probe(list: &str) {
    let coords: Vec<(u32, u32)> = std::fs::read_to_string(list)
        .expect("coords")
        .lines()
        .filter_map(|l| {
            let mut it = l.split_whitespace().map(|v| v.parse().ok());
            Some((it.next()??, it.next()??))
        })
        .collect();
    for line in std::io::stdin().lock().lines() {
        let path = line.expect("stdin");
        let Ok(bytes) = std::fs::read(&path) else {
            continue;
        };
        let (Ok(raw), Ok(meta)) = (dr_decode::decode(&bytes), dr_decode::metadata(&bytes)) else {
            continue;
        };
        let w = raw.width as i64;
        let at = |x: i64, y: i64| {
            let cell = (((y - raw.crop.y as i64) & 1) * 2 + ((x - raw.crop.x as i64) & 1)) as usize;
            raw.data[(y * w + x) as usize].saturating_sub(raw.black_level[cell])
        };
        let rows: Vec<String> = coords
            .iter()
            .map(|&(x, y)| {
                let (x, y) = (x as i64, y as i64);
                let mut n: Vec<u16> = Vec::new();
                for dy in [-2i64, 0, 2] {
                    for dx in [-2i64, 0, 2] {
                        if (dx, dy) != (0, 0) {
                            n.push(at(x + dx, y + dy));
                        }
                    }
                }
                n.sort_unstable();
                format!("[{},{},{}]", at(x, y), n[n.len() - 1], n[n.len() / 2])
            })
            .collect();
        println!(
            "{{\"path\":{:?},\"captured\":{},\"iso\":{},\"shutter\":{},\"range\":{},\"p\":[{}]}}",
            path,
            meta.captured_at.unwrap_or(0),
            meta.iso.unwrap_or(0),
            meta.shutter.unwrap_or(0.0),
            raw.white_level - raw.black_level[0],
            rows.join(","),
        );
    }
}