From 7d5b68c639903c0f597aeea90d3fd290097330bd Mon Sep 17 00:00:00 2001 From: Nathan Vegdahl Date: Thu, 14 Apr 2022 13:34:17 -0700 Subject: [PATCH] Implement deriving transfer function LUTs from processed test images. --- Cargo.lock | 62 ++++++++++++++++++++ Cargo.toml | 4 +- src/main.rs | 165 +++++++++++++++++++++++++++++++++++++++++++++++++--- 3 files changed, 220 insertions(+), 11 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index a15765c..d0cddfa 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -20,6 +20,12 @@ version = "0.10.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "dcb6dd1c2376d2e096796e234a70e17e94cc2d5d54ff8ce42b28cef1d0d359a4" +[[package]] +name = "bitflags" +version = "1.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bef38d45163c2f1dde094a7dfd33ccf595c92905c8f8f4fdc18d06fb1037718a" + [[package]] name = "bumpalo" version = "3.9.1" @@ -32,6 +38,23 @@ version = "1.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "baf1de4339761588bc0619e3cbc0120ee582ebb74b53b4efbf79117bd2da40fd" +[[package]] +name = "clap" +version = "3.1.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "71c47df61d9e16dc010b55dba1952a57d8c215dbb533fd13cdd13369aac73b1c" +dependencies = [ + "bitflags", + "indexmap", + "os_str_bytes", + "textwrap", +] + +[[package]] +name = "colorbox" +version = "0.2.0" +source = "git+https://github.com/cessen/colorbox?branch=master#45d5bf1b747d181fd5db39ba8e80f1f801cc18b2" + [[package]] name = "deflate" version = "1.0.0" @@ -101,6 +124,12 @@ version = "1.8.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "eabb4a44450da02c90444cf74558da904edde8fb4e9035a9a6a4e15445af0bd7" +[[package]] +name = "hashbrown" +version = "0.11.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ab5ef0d4909ef3724cc8cce6ccc8572c5c817592e9285f5464f8e86f8bd3726e" + [[package]] name = "hermit-abi" version = "0.1.19" @@ -110,6 +139,16 @@ dependencies = [ "libc", ] +[[package]] +name = "indexmap" +version = "1.8.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0f647032dfaa1f8b6dc29bd3edb7bbef4861b8b8007ebb118d6db284fd59f6ee" +dependencies = [ + "autocfg", + "hashbrown", +] + [[package]] name = "inflate" version = "0.4.5" @@ -169,9 +208,17 @@ dependencies = [ name = "lut_extractor" version = "0.1.0" dependencies = [ + "clap", + "colorbox", "exr", ] +[[package]] +name = "memchr" +version = "2.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "308cc39be01b73d0d18f82a0e7b2a3df85245f84af96fdddc5d202d27e47b86a" + [[package]] name = "nanorand" version = "0.7.0" @@ -191,6 +238,15 @@ dependencies = [ "libc", ] +[[package]] +name = "os_str_bytes" +version = "6.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8e22443d1643a904602595ba1cd8f7d896afe56d26712531c5ff73a15b2fbf64" +dependencies = [ + "memchr", +] + [[package]] name = "pin-project" version = "1.0.10" @@ -261,6 +317,12 @@ dependencies = [ "unicode-xid", ] +[[package]] +name = "textwrap" +version = "0.15.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b1141d4d61095b28419e22cb0bbf02755f5e54e0526f97f1e3d1d160e60885fb" + [[package]] name = "threadpool" version = "1.8.1" diff --git a/Cargo.toml b/Cargo.toml index 013f441..fd38698 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -3,7 +3,7 @@ name = "lut_extractor" version = "0.1.0" edition = "2021" -# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html - [dependencies] exr = "1.4.1" +clap = { version = "3.1.8", default-features = false, features=["std"] } +colorbox = { git = "https://github.com/cessen/colorbox", branch = "master" } diff --git a/src/main.rs b/src/main.rs index a7663ab..06db9f0 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,27 +1,174 @@ mod test_image; -use test_image::{RES_X, RES_Y}; +use std::{fs::File, io::BufWriter, path::Path}; + +use clap::{Arg, Command}; +use colorbox::transfer_functions::srgb; + +use test_image::{GRADIENT_LEN, RES_X, RES_Y}; + +const VERSION: &str = env!("CARGO_PKG_VERSION"); +const LUT_LEN: usize = 1 << 17; fn main() { - // Build the test image. - let pixels = test_image::build(); + let args = Command::new("LUT Extractor") + .version(VERSION) + .about("A small utility for extracting color LUTs from other software.") + .arg( + Arg::new("input") + .short('i') + .long("input") + .value_name("FILE") + .help("A processed test EXR file, to extract a LUT.") + .takes_value(true) + .required_unless_present_any(&["test_image"]), + ) + .arg( + Arg::new("test_image") + .short('t') + .long("test_image") + .help("Generates the test EXR file."), + ) + .get_matches(); - // Write the test image. + if args.is_present("test_image") { + // Build the test image. + let pixels = test_image::build(); + + // Write the test image. + write_rgb_exr( + &format!("lut_extractor_{}x{}.exr", RES_X, RES_Y), + &pixels, + RES_X, + RES_Y, + ) + .unwrap(); + } else { + let input_path = args.value_of("input").unwrap(); + + let base_image = test_image::build(); + let mut input_image = { + let (image, res_x, res_y) = read_rgb_exr(input_path).unwrap(); + assert_eq!( + res_x, RES_X, + "Input image doesn't have the correct resolution." + ); + assert_eq!( + res_y, RES_Y, + "Input image doesn't have the correct resolution." + ); + image + }; + + // Build the LUT. + let gray = &mut input_image[test_image::gray_idx(0)..test_image::gray_idx(GRADIENT_LEN)]; + let mut prev = gray[0]; + for rgb in gray.iter_mut() { + // Ensure montonicity. + if rgb[0] < prev[0] { + rgb[0] = prev[0]; + } + if rgb[1] < prev[1] { + rgb[1] = prev[1]; + } + if rgb[2] < prev[2] { + rgb[2] = prev[2]; + } + prev = *rgb; + } + let mut gray_r = Vec::with_capacity(LUT_LEN); + let mut gray_g = Vec::with_capacity(LUT_LEN); + let mut gray_b = Vec::with_capacity(LUT_LEN); + for i in 0..LUT_LEN { + let t = i as f32 / (LUT_LEN - 1) as f32; + let rgb = lerp_slice(gray, t); + gray_r.push(rgb[0]); + gray_g.push(rgb[1]); + gray_b.push(rgb[2]); + // gray_r.push(srgb::to_linear(rgb[0])); + // gray_g.push(srgb::to_linear(rgb[1])); + // gray_b.push(srgb::to_linear(rgb[2])); + } + + // Write the LUT. + colorbox::formats::cube::write_1d( + BufWriter::new(File::create("test.cube").unwrap()), + [(0.0, 1.0); 3], + [&gray_r, &gray_g, &gray_b], + ) + .unwrap(); + } +} + +fn lerp_slice(slice: &[[f32; 3]], t: f32) -> [f32; 3] { + assert!(!slice.is_empty()); + + let t2 = (slice.len() - 1) as f32 * t; + let t2i = t2 as usize; + + if t2i == (slice.len() - 1) { + *slice.last().unwrap() + } else { + let alpha = t2.fract(); + let inv_alpha = 1.0 - alpha; + let a = slice[t2i]; + let b = slice[t2i + 1]; + [ + (a[0] * inv_alpha) + (b[0] * alpha), + (a[1] * inv_alpha) + (b[1] * alpha), + (a[2] * inv_alpha) + (b[2] * alpha), + ] + } +} + +fn read_rgb_exr>(path: P) -> exr::error::Result<(Vec<[f32; 3]>, usize, usize)> { + use exr::prelude::{ReadChannels, ReadLayers}; + + let image = exr::image::read::read() + .no_deep_data() + .largest_resolution_level() + .rgb_channels( + |res, _| (vec![[0.0f32; 3]; res.0 * res.1], res.0, res.1), + |(pixels, res_x, _res_y), co, (r, g, b): (f32, f32, f32)| { + pixels[co.1 * *res_x + co.0] = [r, g, b]; + }, + ) + .first_valid_layer() + .all_attributes() + .from_file(path)?; + + Ok(image.layer_data.channel_data.pixels) +} + +fn write_rgb_exr>( + path: P, + pixels: &[[f32; 3]], + res_x: usize, + res_y: usize, +) -> std::io::Result<()> { use exr::{ image::{Encoding, Image, Layer, SpecificChannels}, meta::header::LayerAttributes, prelude::WritableImage, }; - Image::from_layer(Layer::new( - (RES_X, RES_Y), + + assert_eq!(pixels.len(), res_x * res_y); + + match Image::from_layer(Layer::new( + (res_x, res_y), LayerAttributes::named(""), Encoding::SMALL_LOSSLESS, SpecificChannels::rgb(|co: exr::math::Vec2| { - let rgb = pixels[co.1 * RES_X + co.0]; + let rgb = pixels[co.1 * res_y + co.0]; (rgb[0], rgb[1], rgb[2]) }), )) .write() - .to_file(format!("lut_extractor_{}x{}.exr", RES_X, RES_Y)) - .unwrap(); + .to_file(path) + { + // Only IO errors should be possible here. + Err(exr::error::Error::Io(e)) => Err(e), + Err(_) => panic!(), + Ok(()) => Ok(()), + } }