feat: transparency

src/bin.rs

@@ -12,7 +12,8 @@ use clap::{
 };
 use clap_complete::{generate, Shell};
 use dirs::cache_dir;
-use lutgen::{identity, interpolation::*, GenerateLut, Image};
+use image::DynamicImage;
+use lutgen::{identity, interpolation::*, GenerateLut, LutImage};
 use lutgen_palettes::Palette;
 use spinners::{Spinner, Spinners};
 
@@ -60,6 +61,9 @@ enum Subcommands {
             conflicts_with = "force"
         )]
         hald_clut: Option<PathBuf>,
+        /// Enable image transparency
+        #[arg(short, long, default_value_t = false)]
+        transparency: bool,
         /// Enable caching the generated LUT
         #[arg(short, long, default_value_t = false)]
         cache: bool,
@@ -182,7 +186,7 @@ enum Algorithm {
 }
 
 impl LutArgs {
-    fn generate(&self) -> Image {
+    fn generate(&self) -> LutImage {
         let name = self.name();
         let mut sp = Spinner::new(Spinners::Dots3, format!("Generating \"{name}\" LUT..."));
         let time = Instant::now();
@@ -313,7 +317,7 @@ fn main() {
                     lut_args.name(),
                     lut_args.detail_string(),
                 ))),
-                &lut_args.generate(),
+                &lut_args.generate().into(),
             );
 
             println!("Finished in {:?}", total_time.elapsed());
@@ -324,14 +328,15 @@ fn main() {
             lut_args,
             hald_clut,
             images,
+            transparency,
             cache,
             force,
         } => {
             let colors = lut_args.collect();
             // load or generate the lut
             let (hald_clut, details) = {
                 match hald_clut {
-                    Some(path) => (load_image(path), "custom".into()),
+                    Some(path) => (load_image(path).to_rgb8(), "custom".into()),
                     None => {
                         let cache_name =
                             format!("{}_{}", lut_args.name(), lut_args.detail_string());
@@ -345,13 +350,13 @@ fn main() {
 
                             let path = path.join(&cache_name).with_extension("png");
                             if path.exists() && !force {
-                                (load_image(path), cache_name)
+                                (load_image(path).to_rgb8(), cache_name)
                             } else {
                                 if colors.is_empty() {
                                     min_colors_error()
                                 }
                                 let lut = lut_args.generate();
-                                cache_image(path, &lut);
+                                cache_image(path, &lut.clone().into());
                                 (lut, cache_name)
                             }
                         } else {
@@ -365,19 +370,6 @@ fn main() {
             };
 
             for image_path in &images {
-                let mut image_buf = load_image(image_path);
-
-                let mut sp = Spinner::new(
-                    Spinners::Dots3,
-                    format!("Applying LUT to {image_path:?}..."),
-                );
-                let time = Instant::now();
-                identity::correct_image(&mut image_buf, &hald_clut);
-                sp.stop_and_persist(
-                    "✔",
-                    format!("Applied LUT to {image_path:?} in {:?}", time.elapsed()),
-                );
-
                 let output = match images.len() {
                     1 => output.clone().unwrap_or(PathBuf::from(format!(
                         "{}_{details}.png",
@@ -394,7 +386,31 @@ fn main() {
                         folder.join(image_path.file_name().unwrap())
                     },
                 };
-                save_image(output, &image_buf);
+
+                let image = load_image(image_path);
+                let mut sp = Spinner::new(
+                    Spinners::Dots3,
+                    format!("Applying LUT to {image_path:?}..."),
+                );
+                let time = Instant::now();
+
+                if transparency {
+                    let mut image_buf = image.to_rgba8();
+                    identity::correct_image(&mut image_buf, &hald_clut);
+                    sp.stop_and_persist(
+                        "✔",
+                        format!("Applied LUT to {image_path:?} in {:?}", time.elapsed()),
+                    );
+                    save_image(output, &image_buf.into());
+                } else {
+                    let mut image_buf = image.to_rgb8();
+                    identity::correct_image(&mut image_buf, &hald_clut);
+                    sp.stop_and_persist(
+                        "✔",
+                        format!("Applied LUT to {image_path:?} in {:?}", time.elapsed()),
+                    );
+                    save_image(output, &image_buf.into());
+                }
             }
 
             println!("Finished in {:?}", total_time.elapsed());
@@ -410,16 +426,16 @@ fn main() {
     };
 }
 
-fn load_image<P: AsRef<Path>>(path: P) -> Image {
+fn load_image<P: AsRef<Path>>(path: P) -> DynamicImage {
     let path = path.as_ref();
     let mut sp = Spinner::new(Spinners::Dots3, format!("Loading {path:?}..."));
     let time = Instant::now();
-    let lut = image::open(path).expect("failed to open image").to_rgb8();
+    let lut = image::open(path).expect("failed to open image");
     sp.stop_and_persist("✔", format!("Loaded {path:?} in {:?}", time.elapsed()));
     lut
 }
 
-fn save_image<P: AsRef<Path>>(path: P, image: &Image) {
+fn save_image<P: AsRef<Path>>(path: P, image: &DynamicImage) {
     let path = path.as_ref();
     let mut sp = Spinner::new(Spinners::Dots3, format!("Saving output to {path:?}..."));
     let time = Instant::now();
@@ -430,7 +446,7 @@ fn save_image<P: AsRef<Path>>(path: P, image: &Image) {
     );
 }
 
-fn cache_image<P: AsRef<Path>>(path: P, image: &Image) {
+fn cache_image<P: AsRef<Path>>(path: P, image: &DynamicImage) {
     let path = path.as_ref();
     let mut sp = Spinner::new(Spinners::Dots3, format!("Caching {path:?}..."));
     let time = Instant::now();

src/identity.rs

@@ -1,12 +1,12 @@
 //! Hald clut identity creation and application
 
-use image::{ImageBuffer, Rgb};
+use image::{ImageBuffer, Pixel};
 
-use crate::Image;
+use crate::{Image, LutImage};
 
 /// Hald clut base identity generator.
 /// Algorithm derived from: <https://www.quelsolaar.com/technology/clut.html>
-pub fn generate(level: u8) -> Image {
+pub fn generate(level: u8) -> LutImage {
     let level = level as u32;
     let cube_size = level * level;
     let image_size = cube_size * level;
@@ -38,26 +38,34 @@ pub fn generate(level: u8) -> Image {
 ///
 /// Simple implementation that doesn't do any interpolation,
 /// so higher LUT sizes will prove to be more accurate.
-pub fn correct_pixel(input: &[u8; 3], hald_clut: &Image, level: u8) -> [u8; 3] {
+pub fn correct_pixel<P: Pixel<Subpixel = u8>>(input: &mut P, hald_clut: &LutImage, level: u8) {
     let level = level as u32;
     let cube_size = level * level;
 
-    let r = input[0] as u32 * (cube_size - 1) / 255;
-    let g = input[1] as u32 * (cube_size - 1) / 255;
-    let b = input[2] as u32 * (cube_size - 1) / 255;
+    let [r, g, b, ..] = input.channels_mut() else {
+        panic!("pixel must have 3 channels")
+    };
 
-    let x = (r % cube_size) + (g % level) * cube_size;
-    let y = (b * level) + (g / level);
+    let rs = *r as u32 * (cube_size - 1) / 255;
+    let gs = *g as u32 * (cube_size - 1) / 255;
+    let bs = *b as u32 * (cube_size - 1) / 255;
 
-    hald_clut.get_pixel(x, y).0
+    let x = (rs % cube_size) + (gs % level) * cube_size;
+    let y = (bs * level) + (gs / level);
+
+    [*r, *g, *b] = hald_clut.get_pixel(x, y).0;
 }
 
 /// Correct an image with a hald clut identity in place.
 /// Panics if the hald clut is invalid.
 ///
 /// Simple implementation that doesn't do any interpolation,
 /// so higher LUT sizes will prove to be more accurate.
-pub fn correct_image(image: &mut Image, hald_clut: &Image) {
+///
+/// # Panics
+///
+/// If the hald clut is not a square or valid size for the level
+pub fn correct_image<P: Pixel<Subpixel = u8>>(image: &mut Image<P>, hald_clut: &LutImage) {
     let (width, height) = hald_clut.dimensions();
 
     // Find the smallest level that fits inside the hald clut
@@ -72,6 +80,6 @@ pub fn correct_image(image: &mut Image, hald_clut: &Image) {
 
     // Correct the original pixels
     for pixel in image.pixels_mut() {
-        *pixel = Rgb(correct_pixel(&pixel.0, hald_clut, level as u8));
+        correct_pixel(pixel, hald_clut, level as u8);
     }
 }

src/interpolation/mod.rs

@@ -7,7 +7,7 @@ pub use nearest_neighbor::NearestNeighborRemapper;
 use rayon::prelude::{IntoParallelRefMutIterator, ParallelIterator};
 pub use rbf::{gaussian::GaussianRemapper, linear::LinearRemapper, shepard::ShepardRemapper};
 
-use crate::Image;
+use crate::LutImage;
 
 mod gaussian_sample;
 mod nearest_neighbor;
@@ -20,7 +20,7 @@ pub trait InterpolatedRemapper<'a>: Sync {
 
     /// Remap an image in place. Default implementation uses `rayon` to iterate in parallel over
     /// the pixels.
-    fn remap_image(&self, image: &mut Image) {
+    fn remap_image(&self, image: &mut LutImage) {
         image
             .pixels_mut()
             .collect::<Vec<_>>()

src/interpolation/nearest_neighbor.rs

@@ -2,7 +2,7 @@ use image::Rgb;
 use oklab::{srgb_to_oklab, Oklab};
 
 use super::{squared_euclidean, ColorTree, InterpolatedRemapper};
-use crate::{GenerateLut, Image};
+use crate::{GenerateLut, LutImage};
 
 /// Simple remapper that doesn't do any interpolation. Mostly used internally by the other
 /// algorithms.
@@ -30,7 +30,7 @@ impl<'a> NearestNeighborRemapper<'a> {
 
 impl<'a> GenerateLut<'a> for NearestNeighborRemapper<'a> {}
 impl<'a> InterpolatedRemapper<'a> for NearestNeighborRemapper<'a> {
-    fn remap_image(&self, image: &mut Image) {
+    fn remap_image(&self, image: &mut LutImage) {
         for pixel in image.pixels_mut() {
             self.remap_pixel(pixel)
         }

src/lib.rs

@@ -1,17 +1,18 @@
 #![doc = include_str!("../README.md")]
 
-use image::{ImageBuffer, Rgb};
+use image::{ImageBuffer, Pixel, Rgb};
 use interpolation::InterpolatedRemapper;
 
 pub mod identity;
 pub mod interpolation;
 
 /// Core image type (Rgb8)
-pub type Image = ImageBuffer<Rgb<u8>, Vec<u8>>;
+pub type Image<P> = ImageBuffer<P, Vec<<P as Pixel>::Subpixel>>;
+pub type LutImage = Image<Rgb<u8>>;
 
 pub trait GenerateLut<'a>: InterpolatedRemapper<'a> {
     /// Helper method to generate a lut using an [`InterpolatedRemapper`].
-    fn generate_lut(&self, level: u8) -> Image {
+    fn generate_lut(&self, level: u8) -> LutImage {
         let mut identity = identity::generate(level);
         self.remap_image(&mut identity);
         identity