Implement deriving transfer function LUTs from processed test images.

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"

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" }

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() {
+    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();
+
+    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<P: AsRef<Path>>(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<P: AsRef<Path>>(
+    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<usize>| {
-            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))
+    .to_file(path)
-    .unwrap();
+    {
+        // Only IO errors should be possible here.
+        Err(exr::error::Error::Io(e)) => Err(e),
+        Err(_) => panic!(),
+        Ok(()) => Ok(()),
+    }
 }