diff options
| -rw-r--r-- | squash/src/main.rs | 6 | ||||
| -rw-r--r-- | src/difference.rs | 54 | ||||
| -rw-r--r-- | src/lib.rs | 144 |
3 files changed, 134 insertions, 70 deletions
diff --git a/squash/src/main.rs b/squash/src/main.rs index 5437ea1..8e3a79b 100644 --- a/squash/src/main.rs +++ b/squash/src/main.rs @@ -1,5 +1,5 @@ use cli::DifferenceFn; -use colorsquash::{Squasher, SquasherBuilder}; +use colorsquash::SquasherBuilder; use crate::cli::{InType, OutType}; @@ -23,8 +23,8 @@ fn main() -> Result<(), anyhow::Error> { } builder = match cli.difference { - DifferenceFn::Rgb => builder.difference(&colorsquash::difference::rgb_difference), - DifferenceFn::Redmean => builder.difference(&colorsquash::difference::redmean_difference), + DifferenceFn::Rgb => builder.difference(&colorsquash::difference::rgb), + DifferenceFn::Redmean => builder.difference(&colorsquash::difference::redmean), }; let mut squasher = builder.build(&image.data); diff --git a/src/difference.rs b/src/difference.rs index 42c4c4f..b52b34f 100644 --- a/src/difference.rs +++ b/src/difference.rs @@ -1,30 +1,56 @@ -//! A set of difference functions you can use with [SquasherBuilder::difference] +//! A set of difference functions you can use with [SquasherBuilder::difference()] +//! +//! # Writing your own difference function +//! The type you want is `dyn Fn(&RGB8, &RGB8) -> f32` +//! (defined as [`DiffFn`]) +//! +//! The first argument is the color already in the palette and the second is +//! the color we're checking. These are [RGB8] which is a rexport from the `rgb` +//! crate. +//! +//! The value returned is between 0 and 768, but that's not a hard-rule. If you +//! return a value out of that range you'll have to adjust the tolerance with +//! [Squasher::set_tolerance()] or [SquasherBuilder::tolerance]. +//! +//! The difference functions have the possibility of being called hundreds of +//! thousands of times; you might want to `#[inline(always)]` + +// This is used in the module level documentation just above. Without it we'd +// have to fully qualify the interlink which is also how it'd be displayed. +#[allow(unused_imports)] +use crate::{Squasher, SquasherBuilder}; // rexport this so people don't need to add the rgb crate to their project. this // also helps avoid crate version mismatch +/// rexport from the [`rgb`](https://docs.rs/rgb/0.8.37/rgb/) crate. pub use rgb::RGB8; +/// Type definition for difference functions. +pub type DiffFn = dyn Fn(&RGB8, &RGB8) -> f32; + /// A naïve comparison just summing the channel differences /// I.E. `|a.red - b.red| + |a.green - b.green| + |a.blue - b.blue|` #[allow(clippy::many_single_char_names)] #[inline(always)] -pub fn rgb_difference(a: &RGB8, b: &RGB8) -> f32 { - let absdiff = |a: u8, b: u8| (a as f32 - b as f32).abs(); - absdiff(a.r, b.r) + absdiff(a.g, b.g) + absdiff(a.b, b.b) +pub fn rgb(a: &RGB8, b: &RGB8) -> f32 { + let absdiff = |a: u8, b: u8| (a as f32 - b as f32).abs(); + absdiff(a.r, b.r) + absdiff(a.g, b.g) + absdiff(a.b, b.b) } // https://en.wikipedia.org/wiki/Color_difference#sRGB +/// a slightly more intelligent algorithm that weighs the channels in an attempt +/// to better align with human color perception. #[inline(always)] -pub fn redmean_difference(a: &RGB8, b: &RGB8) -> f32 { - let delta_r = a.r as f32 - b.r as f32; - let delta_g = a.g as f32 - b.g as f32; - let delta_b = a.b as f32 - b.b as f32; - // reasonably sure calling it prime is wrong, but - let r_prime = 0.5 * (a.r as f32 + b.r as f32); +pub fn redmean(a: &RGB8, b: &RGB8) -> f32 { + let delta_r = a.r as f32 - b.r as f32; + let delta_g = a.g as f32 - b.g as f32; + let delta_b = a.b as f32 - b.b as f32; + // reasonably sure calling it prime is wrong, but + let r_prime = 0.5 * (a.r as f32 + b.r as f32); - let red_part = (2.0 + (r_prime / 256.0)) * (delta_r * delta_r); - let green_part = 4.0 * (delta_g * delta_g); - let blue_part = (2.0 + (255.0 - r_prime) / 256.0) * (delta_b * delta_b); + let red_part = (2.0 + (r_prime / 256.0)) * (delta_r * delta_r); + let green_part = 4.0 * (delta_g * delta_g); + let blue_part = (2.0 + (255.0 - r_prime) / 256.0) * (delta_b * delta_b); - (red_part + green_part + blue_part).sqrt() + (red_part + green_part + blue_part).sqrt() } diff --git a/src/lib.rs b/src/lib.rs index 40a8f28..38f8e1a 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,11 +1,11 @@ #[cfg(kmeans)] use kmeans::{KMeans, KMeansConfig}; -use rgb::RGB8; +use rgb::{ComponentBytes, FromSlice, RGB8}; use std::collections::HashMap; pub mod difference; -type DiffFn = dyn Fn(&RGB8, &RGB8) -> f32; +use difference::DiffFn; pub struct SquasherBuilder<T> { max_colours: T, @@ -28,7 +28,8 @@ impl<T: Count> SquasherBuilder<T> { /// The function to use to compare colours. /// - /// see the [difference] module for functions included with the crate. + /// see the [difference] module for functions included with the crate and + /// information on implementing your own. pub fn difference(mut self, difference: &'static DiffFn) -> SquasherBuilder<T> { self.difference_fn = Box::new(difference); self @@ -41,7 +42,10 @@ impl<T: Count> SquasherBuilder<T> { self } - pub fn build(self, image: &[u8]) -> Squasher<T> { + pub fn build<'a, Img>(self, image: Img) -> Squasher<T> + where + Img: Into<ImageData<'a>>, + { let mut squasher = Squasher::from_parts(self.max_colours, self.difference_fn, self.tolerance); squasher.recolor(image); @@ -54,7 +58,7 @@ impl<T: Count> Default for SquasherBuilder<T> { fn default() -> Self { Self { max_colours: T::from_usize(255), - difference_fn: Box::new(difference::rgb_difference), + difference_fn: Box::new(difference::rgb), tolerance: 1.0, } } @@ -73,12 +77,11 @@ impl<T: Count> Squasher<T> { /// Creates a new squasher and allocates a new color map. A color map /// contains every 24-bit color and ends up with an amount of memory /// equal to `16MB * std::mem::size_of(T)`. - pub fn new(max_colors_minus_one: T, buffer: &[u8]) -> Self { - let mut this = Self::from_parts( - max_colors_minus_one, - Box::new(difference::rgb_difference), - 1.0, - ); + pub fn new<'a, Img>(max_colors_minus_one: T, buffer: Img) -> Self + where + Img: Into<ImageData<'a>>, + { + let mut this = Self::from_parts(max_colors_minus_one, Box::new(difference::rgb), 1.0); this.recolor(buffer); this @@ -88,23 +91,35 @@ impl<T: Count> Squasher<T> { SquasherBuilder::new() } + /// Set the tolerance pub fn set_tolerance(&mut self, percent: f32) { self.tolerance_percent = percent; } /// Create a new palette from the colours in the given image. #[cfg(not(kmeans))] - pub fn recolor(&mut self, image: &[u8]) { + pub fn recolor<'a, Img>(&mut self, image: Img) + where + Img: Into<ImageData<'a>>, + { let sorted = Self::unique_and_sort(image); let selected = self.select_colors(sorted); self.palette = selected; } #[cfg(kmeans)] - pub fn recolor(&mut self, image: &[u8]) { + pub fn recolor<'a, Img>(&mut self, image: Img) + where + Img: Into<ImageData<'a>>, + { + let ImageData(rgb) = image.into(); + let kmean = KMeans::new( - image.iter().map(|u| *u as f32).collect::<Vec<f32>>(), - image.len() / 3, + rgb.as_bytes() + .iter() + .map(|u| *u as f32) + .collect::<Vec<f32>>(), + rgb.as_bytes().len() / 3, 3, ); let k = self.max_colours_min1.as_usize() + 1; @@ -137,43 +152,47 @@ impl<T: Count> Squasher<T> { /// Take an RGB image buffer and an output buffer. The function will fill /// the output buffer with indexes into the Palette. The output buffer should /// be a third of the size of the image buffer. - pub fn map(&mut self, image: &[u8], buffer: &mut [T]) { - if buffer.len() * 3 < image.len() { - panic!("outout buffer too small to fit indexed image"); + pub fn map<'a, Img>(&mut self, image: Img, buffer: &mut [T]) + where + Img: Into<ImageData<'a>>, + { + let ImageData(rgb) = image.into(); + + if buffer.len() * 3 < rgb.len() { + panic!("output buffer too small to fit indexed image"); } // We have to map the colours of this image now because it might contain // colours not present in the first image. - let sorted = Self::unique_and_sort(image); + let sorted = Self::unique_and_sort(rgb); self.map_selected(&sorted); - for (idx, color) in image.chunks(3).enumerate() { - let index = self.map[color_index(&RGB8::new(color[0], color[1], color[2]))]; - - buffer[idx] = index; + for (idx, color) in rgb.iter().enumerate() { + buffer[idx] = self.map[color_index(color)]; } } /// Like [Squasher::map] but it doesn't recount the input image. This will /// cause colors the Squasher hasn't seen before to come out as index 0 which - /// may be incorrect. + /// may be incorrect! //TODO: gen- Better name? - pub fn map_unsafe(&self, image: &[u8], buffer: &mut [T]) { - if buffer.len() * 3 < image.len() { - panic!("outout buffer too small to fit indexed image"); + pub fn map_no_recolor<'a, Img>(&self, image: Img, buffer: &mut [T]) + where + Img: Into<ImageData<'a>>, + { + let ImageData(rgb) = image.into(); + + if buffer.len() * 3 < rgb.len() { + panic!("output buffer too small to fit indexed image"); } - for (idx, color) in image.chunks(3).enumerate() { - let index = self.map[color_index(&RGB8::new(color[0], color[1], color[2]))]; - - buffer[idx] = index; + for (idx, color) in rgb.iter().enumerate() { + buffer[idx] = self.map[color_index(color)]; } } #[cfg(feature = "gifed")] pub fn palette_gifed(&self) -> gifed::block::Palette { - use rgb::ComponentBytes; - self.palette.as_slice().as_bytes().try_into().unwrap() } @@ -184,24 +203,22 @@ impl<T: Count> Squasher<T> { /// Retrieve the palette as bytes pub fn palette_bytes(&self) -> Vec<u8> { - self.palette - .clone() - .into_iter() - .flat_map(|rgb| [rgb.r, rgb.g, rgb.b].into_iter()) - .collect() + self.palette.as_bytes().to_owned() } /// Takes an image buffer of RGB data and fill the color map - fn unique_and_sort(buffer: &[u8]) -> Vec<RGB8> { + fn unique_and_sort<'a, Img>(buffer: Img) -> Vec<RGB8> + where + Img: Into<ImageData<'a>>, + { + let ImageData(rgb) = buffer.into(); let mut colors: HashMap<RGB8, usize> = HashMap::default(); //count pixels - for pixel in buffer.chunks(3) { - let rgb = RGB8::new(pixel[0], pixel[1], pixel[2]); - - match colors.get_mut(&rgb) { + for px in rgb { + match colors.get_mut(px) { None => { - colors.insert(rgb, 1); + colors.insert(*px, 1); } Some(n) => *n += 1, } @@ -227,10 +244,6 @@ impl<T: Count> Squasher<T> { /// of times they occur, high to low. #[cfg(not(kmeans))] fn select_colors(&self, sorted: Vec<RGB8>) -> Vec<RGB8> { - // I made these numbers up - #[allow(non_snake_case)] - //let RGB_TOLERANCE: f32 = 0.01 * 765.0; - //let RGB_TOLERANCE: f32 = 36.0; let tolerance = (self.tolerance_percent / 100.0) * 765.0; let max_colours = self.max_colours_min1.as_usize() + 1; let mut selected_colors: Vec<RGB8> = Vec::with_capacity(max_colours); @@ -238,10 +251,9 @@ impl<T: Count> Squasher<T> { for sorted_color in sorted { if max_colours <= selected_colors.len() { break; - } else if selected_colors - .iter() - .all(|color| (self.difference_fn)(&sorted_color, color) > tolerance) - { + } else if selected_colors.iter().all(|selected_color| { + (self.difference_fn)(selected_color, &sorted_color) > tolerance + }) { selected_colors.push(sorted_color); } } @@ -276,9 +288,14 @@ impl Squasher<u8> { /// # Returns /// The new size of the image pub fn map_over(&mut self, image: &mut [u8]) -> usize { + // "redundant slicing" here is to drop the mut on the reference because + // ImageData doesn't have a From<&mut [u8]> and I don't particularly want + // it to + #[allow(clippy::redundant_slicing)] + let sorted = Self::unique_and_sort(&image[..]); + // We have to map the colours of this image now because it might contain // colours not present in the first image. - let sorted = Self::unique_and_sort(image); self.map_selected(&sorted); for idx in 0..(image.len() / 3) { @@ -330,6 +347,27 @@ count_impl!(u32); count_impl!(u64); count_impl!(usize); +pub struct ImageData<'a>(&'a [RGB8]); + +impl<'a> From<&'a Vec<u8>> for ImageData<'a> { + fn from(plain: &'a Vec<u8>) -> Self { + ImageData(plain.as_rgb()) + } +} + +impl<'a> From<&'a [u8]> for ImageData<'a> { + fn from(plain: &'a [u8]) -> Self { + ImageData(plain.as_rgb()) + } +} + +impl<'a> From<&'a [RGB8]> for ImageData<'a> { + fn from(rgb: &'a [RGB8]) -> Self { + ImageData(rgb) + } +} + +/// Compute the color index into the big-map-of-all-colours. #[inline(always)] fn color_index(c: &RGB8) -> usize { c.r as usize * (256 * 256) + c.g as usize * 256 + c.b as usize |