dtourolle
d4248bc0dd
The learned demosaic replaces the classical one and takes its input, the mosaic hot_pixels.wgsl leaves (denoise.md §2), so its training data and its input in the app must come through that pass and not a lookalike. The pass was recorded inline in Demosaicer::run; it is now built by hot_pass and recorded by record_hot_pass, which run still uses unchanged, and Demosaicer::repair_hot_pixels runs it on its own and reads the mosaic back. mosaic_dump moves to dr-gpu to call it, records how many photosites changed, and keeps --unrepaired for a raw readout.
116 lines
4.4 KiB
Rust
116 lines
4.4 KiB
Rust
//! Dump RAW files' mosaics for training the learned denoise (FR-DEV-3g).
|
|
//!
|
|
//! The training repo must read photosites the way the app reads them —
|
|
//! same black and white levels, same active area, same CFA phase — or a
|
|
//! network trained on one phase runs on another and paints moiré everywhere
|
|
//! (denoise.md §4.4). So it reads this, not LibRaw.
|
|
//!
|
|
//! The photosites are those the demosaic reads: hot and dead ones repaired by
|
|
//! the app's own pass ([`Demosaicer::repair_hot_pixels`], the same shader
|
|
//! `run` dispatches), because the learned stage replaces the demosaic and
|
|
//! takes its input (denoise.md §2). `--unrepaired` skips it.
|
|
//!
|
|
//! Reads `input<TAB>output-prefix` lines on stdin and writes, per line,
|
|
//! `prefix.npy` (the whole readout, masked border included, `u16`, row-major)
|
|
//! and `prefix.json` (what `decode` and `metadata` say about it). The border
|
|
//! is kept, and the repair never touches it, because its optically black
|
|
//! photosites are a dark frame for free: read noise and row noise at that ISO.
|
|
//!
|
|
//! ```sh
|
|
//! printf 'IMG_0001.CR2\tout/IMG_0001\n' |
|
|
//! cargo run --release -p dr-gpu --example mosaic_dump
|
|
//! ```
|
|
|
|
use std::io::{BufRead, Write};
|
|
|
|
use dr_gpu::{Demosaicer, GpuContext};
|
|
|
|
fn main() {
|
|
let repair = !std::env::args().any(|a| a == "--unrepaired");
|
|
let ctx = pollster::block_on(GpuContext::new_headless()).expect("a GPU for the hot-pixel pass");
|
|
let demosaicer = Demosaicer::new(&ctx).expect("demosaicer");
|
|
let mut failed = 0;
|
|
for line in std::io::stdin().lock().lines() {
|
|
let line = line.expect("stdin");
|
|
let Some((input, prefix)) = line.split_once('\t') else {
|
|
continue;
|
|
};
|
|
match dump(input, prefix, repair.then_some(&demosaicer)) {
|
|
Ok(()) => println!("ok\t{input}"),
|
|
Err(e) => {
|
|
failed += 1;
|
|
println!("fail\t{input}\t{e}");
|
|
}
|
|
}
|
|
std::io::stdout().flush().ok();
|
|
}
|
|
std::process::exit(if failed > 0 { 1 } else { 0 });
|
|
}
|
|
|
|
fn dump(input: &str, prefix: &str, repair: Option<&Demosaicer>) -> Result<(), String> {
|
|
let bytes = std::fs::read(input).map_err(|e| e.to_string())?;
|
|
let mut raw = dr_decode::decode(&bytes).map_err(|e| e.to_string())?;
|
|
if raw.samples_per_pixel != 1 {
|
|
return Err("linear DNG: no photosites".into());
|
|
}
|
|
let repaired = match repair {
|
|
Some(d) => d.repair_hot_pixels(&mut raw).map_err(|e| e.to_string())? as i64,
|
|
None => -1,
|
|
};
|
|
let meta = dr_decode::metadata(&bytes).map_err(|e| e.to_string())?;
|
|
|
|
let mut npy = Vec::with_capacity(raw.data.len() * 2 + 128);
|
|
let mut header = format!(
|
|
"{{'descr': '<u2', 'fortran_order': False, 'shape': ({}, {}), }}",
|
|
raw.height, raw.width
|
|
);
|
|
// The header, its magic and length are padded to a multiple of 64.
|
|
while (10 + header.len() + 1) % 64 != 0 {
|
|
header.push(' ');
|
|
}
|
|
header.push('\n');
|
|
npy.extend_from_slice(b"\x93NUMPY\x01\x00");
|
|
npy.extend_from_slice(&(header.len() as u16).to_le_bytes());
|
|
npy.extend_from_slice(header.as_bytes());
|
|
for v in &raw.data {
|
|
npy.extend_from_slice(&v.to_le_bytes());
|
|
}
|
|
std::fs::write(format!("{prefix}.npy"), npy).map_err(|e| e.to_string())?;
|
|
|
|
let opt = |v: Option<f32>| v.map_or("null".to_string(), |v| v.to_string());
|
|
let matrix = raw
|
|
.color_matrix
|
|
.map_or("null".to_string(), |m| format!("{m:?}"));
|
|
let json = format!(
|
|
concat!(
|
|
"{{\"source\": {:?}, \"make\": {:?}, \"model\": {:?}, ",
|
|
"\"width\": {}, \"height\": {}, ",
|
|
"\"crop\": [{}, {}, {}, {}], \"cfa\": {:?}, ",
|
|
"\"black\": {:?}, \"white\": {}, \"wb\": {:?}, \"cam_to_srgb\": {}, ",
|
|
"\"iso\": {}, \"shutter\": {}, \"aperture\": {}, \"captured_at\": {}, ",
|
|
"\"hot_repaired\": {}}}\n"
|
|
),
|
|
input,
|
|
raw.make,
|
|
raw.model,
|
|
raw.width,
|
|
raw.height,
|
|
raw.crop.x,
|
|
raw.crop.y,
|
|
raw.crop.width,
|
|
raw.crop.height,
|
|
format!("{:?}", raw.cfa_pattern),
|
|
raw.black_level,
|
|
raw.white_level,
|
|
raw.wb_coeffs,
|
|
matrix,
|
|
meta.iso.map_or("null".to_string(), |v| v.to_string()),
|
|
opt(meta.shutter),
|
|
opt(meta.aperture),
|
|
meta.captured_at
|
|
.map_or("null".to_string(), |v| v.to_string()),
|
|
repaired,
|
|
);
|
|
std::fs::write(format!("{prefix}.json"), json).map_err(|e| e.to_string())
|
|
}
|