diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..2af5d3a --- /dev/null +++ b/.gitignore @@ -0,0 +1,3 @@ +/target +/.DS_Store +/.vscode \ No newline at end of file diff --git a/Cargo.lock b/Cargo.lock new file mode 100644 index 0000000..8d514a5 --- /dev/null +++ b/Cargo.lock @@ -0,0 +1,331 @@ +# This file is automatically @generated by Cargo. +# It is not intended for manual editing. +version = 3 + +[[package]] +name = "adler" +version = "1.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f26201604c87b1e01bd3d98f8d5d9a8fcbb815e8cedb41ffccbeb4bf593a35fe" + +[[package]] +name = "adler32" +version = "1.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "aae1277d39aeec15cb388266ecc24b11c80469deae6067e17a1a7aa9e5c1f234" + +[[package]] +name = "autocfg" +version = "1.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cdb031dd78e28731d87d56cc8ffef4a8f36ca26c38fe2de700543e627f8a464a" + +[[package]] +name = "bitflags" +version = "1.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bef38d45163c2f1dde094a7dfd33ccf595c92905c8f8f4fdc18d06fb1037718a" + +[[package]] +name = "bytemuck" +version = "1.7.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "439989e6b8c38d1b6570a384ef1e49c8848128f5a97f3914baef02920842712f" + +[[package]] +name = "byteorder" +version = "1.4.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "14c189c53d098945499cdfa7ecc63567cf3886b3332b312a5b4585d8d3a6a610" + +[[package]] +name = "cfg-if" +version = "1.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "baf1de4339761588bc0619e3cbc0120ee582ebb74b53b4efbf79117bd2da40fd" + +[[package]] +name = "color_quant" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3d7b894f5411737b7867f4827955924d7c254fc9f4d91a6aad6b097804b1018b" + +[[package]] +name = "crc32fast" +version = "1.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "738c290dfaea84fc1ca15ad9c168d083b05a714e1efddd8edaab678dc28d2836" +dependencies = [ + "cfg-if", +] + +[[package]] +name = "crossbeam-channel" +version = "0.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "06ed27e177f16d65f0f0c22a213e17c696ace5dd64b14258b52f9417ccb52db4" +dependencies = [ + "cfg-if", + "crossbeam-utils", +] + +[[package]] +name = "crossbeam-deque" +version = "0.8.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6455c0ca19f0d2fbf751b908d5c55c1f5cbc65e03c4225427254b46890bdde1e" +dependencies = [ + "cfg-if", + "crossbeam-epoch", + "crossbeam-utils", +] + +[[package]] +name = "crossbeam-epoch" +version = "0.9.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4ec02e091aa634e2c3ada4a392989e7c3116673ef0ac5b72232439094d73b7fd" +dependencies = [ + "cfg-if", + "crossbeam-utils", + "lazy_static", + "memoffset", + "scopeguard", +] + +[[package]] +name = "crossbeam-utils" +version = "0.8.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d82cfc11ce7f2c3faef78d8a684447b40d503d9681acebed6cb728d45940c4db" +dependencies = [ + "cfg-if", + "lazy_static", +] + +[[package]] +name = "deflate" +version = "0.8.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "73770f8e1fe7d64df17ca66ad28994a0a623ea497fa69486e14984e715c5d174" +dependencies = [ + "adler32", + "byteorder", +] + +[[package]] +name = "either" +version = "1.6.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e78d4f1cc4ae33bbfc157ed5d5a5ef3bc29227303d595861deb238fcec4e9457" + +[[package]] +name = "gif" +version = "0.11.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c3a7187e78088aead22ceedeee99779455b23fc231fe13ec443f99bb71694e5b" +dependencies = [ + "color_quant", + "weezl", +] + +[[package]] +name = "hermit-abi" +version = "0.1.19" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "62b467343b94ba476dcb2500d242dadbb39557df889310ac77c5d99100aaac33" +dependencies = [ + "libc", +] + +[[package]] +name = "image" +version = "0.23.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "24ffcb7e7244a9bf19d35bf2883b9c080c4ced3c07a9895572178cdb8f13f6a1" +dependencies = [ + "bytemuck", + "byteorder", + "color_quant", + "gif", + "jpeg-decoder", + "num-iter", + "num-rational", + "num-traits", + "png", + "scoped_threadpool", + "tiff", +] + +[[package]] +name = "jpeg-decoder" +version = "0.1.22" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "229d53d58899083193af11e15917b5640cd40b29ff475a1fe4ef725deb02d0f2" +dependencies = [ + "rayon", +] + +[[package]] +name = "lazy_static" +version = "1.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e2abad23fbc42b3700f2f279844dc832adb2b2eb069b2df918f455c4e18cc646" + +[[package]] +name = "libc" +version = "0.2.112" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1b03d17f364a3a042d5e5d46b053bbbf82c92c9430c592dd4c064dc6ee997125" + +[[package]] +name = "memoffset" +version = "0.6.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5aa361d4faea93603064a027415f07bd8e1d5c88c9fbf68bf56a285428fd79ce" +dependencies = [ + "autocfg", +] + +[[package]] +name = "miniz_oxide" +version = "0.3.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "791daaae1ed6889560f8c4359194f56648355540573244a5448a83ba1ecc7435" +dependencies = [ + "adler32", +] + +[[package]] +name = "miniz_oxide" +version = "0.4.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a92518e98c078586bc6c934028adcca4c92a53d6a958196de835170a01d84e4b" +dependencies = [ + "adler", + "autocfg", +] + +[[package]] +name = "num-integer" +version = "0.1.44" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d2cc698a63b549a70bc047073d2949cce27cd1c7b0a4a862d08a8031bc2801db" +dependencies = [ + "autocfg", + "num-traits", +] + +[[package]] +name = "num-iter" +version = "0.1.42" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b2021c8337a54d21aca0d59a92577a029af9431cb59b909b03252b9c164fad59" +dependencies = [ + "autocfg", + "num-integer", + "num-traits", +] + +[[package]] +name = "num-rational" +version = "0.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "12ac428b1cb17fce6f731001d307d351ec70a6d202fc2e60f7d4c5e42d8f4f07" +dependencies = [ + "autocfg", + "num-integer", + "num-traits", +] + +[[package]] +name = "num-traits" +version = "0.2.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9a64b1ec5cda2586e284722486d802acf1f7dbdc623e2bfc57e65ca1cd099290" +dependencies = [ + "autocfg", +] + +[[package]] +name = "num_cpus" +version = "1.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "05499f3756671c15885fee9034446956fff3f243d6077b91e5767df161f766b3" +dependencies = [ + "hermit-abi", + "libc", +] + +[[package]] +name = "octo" +version = "0.1.0" +dependencies = [ + "image", +] + +[[package]] +name = "png" +version = "0.16.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3c3287920cb847dee3de33d301c463fba14dda99db24214ddf93f83d3021f4c6" +dependencies = [ + "bitflags", + "crc32fast", + "deflate", + "miniz_oxide 0.3.7", +] + +[[package]] +name = "rayon" +version = "1.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c06aca804d41dbc8ba42dfd964f0d01334eceb64314b9ecf7c5fad5188a06d90" +dependencies = [ + "autocfg", + "crossbeam-deque", + "either", + "rayon-core", +] + +[[package]] +name = "rayon-core" +version = "1.9.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d78120e2c850279833f1dd3582f730c4ab53ed95aeaaaa862a2a5c71b1656d8e" +dependencies = [ + "crossbeam-channel", + "crossbeam-deque", + "crossbeam-utils", + "lazy_static", + "num_cpus", +] + +[[package]] +name = "scoped_threadpool" +version = "0.1.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1d51f5df5af43ab3f1360b429fa5e0152ac5ce8c0bd6485cae490332e96846a8" + +[[package]] +name = "scopeguard" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d29ab0c6d3fc0ee92fe66e2d99f700eab17a8d57d1c1d3b748380fb20baa78cd" + +[[package]] +name = "tiff" +version = "0.6.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9a53f4706d65497df0c4349241deddf35f84cee19c87ed86ea8ca590f4464437" +dependencies = [ + "jpeg-decoder", + "miniz_oxide 0.4.4", + "weezl", +] + +[[package]] +name = "weezl" +version = "0.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d8b77fdfd5a253be4ab714e4ffa3c49caf146b4de743e97510c0656cf90f1e8e" diff --git a/Cargo.toml b/Cargo.toml new file mode 100644 index 0000000..95af9d1 --- /dev/null +++ b/Cargo.toml @@ -0,0 +1,12 @@ +[package] +name = "octo" +version = "0.1.0" +edition = "2021" + +# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html + +[dependencies] +image="0.23.14" + +[profile.release] +debug = true \ No newline at end of file diff --git a/README.md b/README.md new file mode 100644 index 0000000..d2877f8 --- /dev/null +++ b/README.md @@ -0,0 +1,6 @@ +Does quantization in pure rust on Rgba image, using https://crates.io/crates/image + +algorithm: +https://observablehq.com/@tmcw/octree-color-quantization + +inspired by to https://github.com/objectProfessionals/movieMaps \ No newline at end of file diff --git a/src/lib.rs b/src/lib.rs new file mode 100644 index 0000000..0a9c63b --- /dev/null +++ b/src/lib.rs @@ -0,0 +1,359 @@ +use std::{cell::RefCell, rc::Rc}; + +use image::{Rgba, RgbaImage}; + +const MAX_LEVEL: usize = 5; + +pub fn quantize(image: &RgbaImage, num_colors:usize) -> RgbaImage{ + let mut quantizer=OctTreeQuantizer::new(num_colors); + quantizer.quantize(image) +} + +struct OctTreeQuantizer { + root: Rc>, + reduce_colors: usize, + maximum_colors: usize, + colors: usize, + color_list: Vec>>>>, +} + +impl OctTreeQuantizer { + fn new(num_colors: usize) -> Self { + let mut new_quantizer = Self { + root: Rc::new(RefCell::new(OctTreeNode::new())), + reduce_colors: usize::max(512, num_colors * 2), + maximum_colors: num_colors, + color_list: vec![], + colors: 0, + }; + + for _ in 0..=MAX_LEVEL { + new_quantizer.color_list.push(Vec::new()); + } + new_quantizer + } + + pub fn quantize(&mut self, image: &RgbaImage) -> RgbaImage { + for y in 0..image.height() { + for x in 0..image.width() { + self.insert_color(image.get_pixel(x, y), Rc::clone(&self.root)); + + if self.colors > self.reduce_colors { + self.reduce_tree(self.reduce_colors); + //reduce sets to None and the code below actually removes nodes from the list + for level in &mut self.color_list { + level.retain(|c| c.is_some()); + } + } + } + } + let table = self.build_color_table(); + + let mut out = RgbaImage::new(image.width(), image.height()); + for y in 0..image.height() { + for x in 0..image.width() { + let pixel = image.get_pixel(x, y); + if let Some(index) = self.get_index_for_color(pixel, &self.root) { + let color = table.get(index).unwrap(); + if let Some(color) = color { + out.put_pixel(x, y, *color); + } + } + } + } + out + } + + fn get_index_for_color( + &self, + color: &Rgba, + node: &Rc>, + ) -> Option { + fn get_index_for_color( + quantizer: &OctTreeQuantizer, + color: &Rgba, + level: usize, + node: &Rc>, + ) -> Option { + if level > MAX_LEVEL { + return None; + } + let node = Rc::clone(node); + let index = get_bitmask(color, &level); + + let node_b = node.borrow(); + let child = &node_b.leaf[index]; + + if let Some(child) = child { + let child_b = child.borrow(); + if child_b.is_leaf { + return Some(child_b.index); + } else { + return get_index_for_color(quantizer, color, level + 1, child); + } + } else { + return Some(node_b.index); + } + } + + get_index_for_color(&self, color, 0, node) + } + + fn build_color_table(&mut self) -> Vec>> { + //nested function that is called recursively + fn build_color_table( + quantizer: &mut OctTreeQuantizer, + node: &Rc>, + table: &mut Vec>>, + index: usize, + ) -> usize { + if quantizer.colors > quantizer.maximum_colors { + quantizer.reduce_tree(quantizer.maximum_colors); + } + if node.borrow().is_leaf { + { + let node = node.borrow(); + let count = node.count; + table[index] = Some(Rgba::from([ + (node.total_red / count as u32) as u8, + (node.total_green / count as u32) as u8, + (node.total_blue / count as u32) as u8, + 0xFF, + ])); + } + node.borrow_mut().index = index; + + index + 1 + } else { + let mut result = index; + for i in 0..8 { + // cannot iterate leaf, because that widens the scope of the borrow (of node) + let mut node = node.borrow_mut(); + if let Some(leaf) = &node.leaf[i] { + //could be immutable borrow + let new_index = build_color_table(quantizer, leaf, table, result); + node.index = index; //but also need mutable borrow here + result = new_index; + } + } + + result + } + } + + let mut table: Vec>> = vec![None; self.colors]; + let node = Rc::clone(&self.root); + build_color_table(self, &node, &mut table, 0); + table + } + + fn insert_color(&mut self, rgb: &Rgba, node: Rc>) { + //nested function that is called recursively + fn insert_color( + quantizer: &mut OctTreeQuantizer, + color: &Rgba, + level: usize, + node: Rc>, + ) { + if level > MAX_LEVEL { + return; + } + + let index = get_bitmask(color, &level); + + if node.borrow().leaf[index].is_none() { + let mut child = OctTreeNode::new(); + child.parent = Some(Rc::clone(&node)); + child.p_index = quantizer.color_list[level].len(); + + if level == MAX_LEVEL { + child.is_leaf = true; + child.count = 1; + child.total_red = color[0] as u32; + child.total_green = color[1] as u32; + child.total_blue = color[2] as u32; + child.level = level; + quantizer.colors += 1; + } + + let child = Rc::new(RefCell::new(child)); + quantizer.color_list[level].push(Some(Rc::clone(&child))); + let clone = Rc::clone(&child); + + { + let mut mutnode = node.borrow_mut(); + mutnode.children += 1; + mutnode.is_leaf = false; + mutnode.leaf[index] = Some(child); + } + + if level < MAX_LEVEL { + insert_color(quantizer, color, level + 1, clone); + } else { + return; + } + } else { + if node + .borrow() + .leaf + .get(index) + .unwrap() + .as_ref() + .unwrap() + .borrow() + .is_leaf + { + let mut node = node.borrow_mut(); + let mut child = node + .leaf + .get_mut(index) + .unwrap() + .as_ref() + .unwrap() + .borrow_mut(); + child.count += 1; + child.total_red += color[0] as u32; + child.total_green += color[1] as u32; + child.total_blue += color[2] as u32; + return; + } else { + insert_color( + quantizer, + color, + level + 1, + Rc::clone(&(node.borrow().leaf[index]).as_ref().unwrap()), + ); + } + } + } + + insert_color(self, rgb, 0, node); + } + + fn reduce_tree(&mut self, num_colors: usize) { + // Nested function that is called recursively + fn reduce_tree(quantizer: &mut OctTreeQuantizer, num_colors: usize, level: isize) { + if level < 0 { + return; + } else { + let mut removals = Vec::new(); + let list = &quantizer.color_list[level as usize]; + for node in list { + if let Some(node) = node { + if node.borrow().children > 0 { + for i in 0..8 { + let mut color: Option<(usize, u32, u32, u32, usize)> = None; + + if let Some(child) = node.borrow().leaf.get(i) { + if let Some(child) = child { + let child = child.borrow(); + color = Some(( + child.count, + child.total_red, + child.total_green, + child.total_blue, + child.p_index, + )); + } + } + + // need to mutate node, which conflicts with previous borrow to retrieve the child + if let Some(color) = color { + let mut node = node.borrow_mut(); + node.count += color.0; + node.total_red += color.1; + node.total_green += color.2; + node.total_blue += color.3; + node.leaf[i] = None; + node.children -= 1; + quantizer.colors -= 1; + + removals.push(color.4); //save for further processing outside loop (and borrow of colorlist) + } + } + node.borrow_mut().is_leaf = true; + quantizer.colors += 1; + if quantizer.colors <= num_colors { + return; + } + } + } + } + let color_list = &mut quantizer.color_list[(level as usize + 1)]; + for index in removals { + color_list[index] = None; //set to None here, Option removed later + } + + reduce_tree(quantizer, num_colors, level - 1); + } + } + + // call to nested function + reduce_tree(self, num_colors, (MAX_LEVEL - 1) as isize); + } +} + +struct OctTreeNode { + children: usize, + level: usize, + parent: Option>>, + leaf: Vec>>>, + is_leaf: bool, + count: usize, + total_red: u32, + total_green: u32, + total_blue: u32, + index: usize, + p_index: usize, +} + +impl OctTreeNode { + fn new() -> Self { + Self { + children: 0, + level: 0, + parent: None, + leaf: vec![None; 8], + is_leaf: false, + count: 0, + total_red: 0, + total_green: 0, + total_blue: 0, + index: 0, + p_index: 0, + } + } +} + +fn get_bitmask(color: &Rgba, level: &usize) -> usize { + let bit = 0x80 >> level; + + let mut index = 0; + if (color[0] & bit) != 0 { + index += 4; + } + if (color[1] & bit) != 0 { + index += 2; + } + if (color[2] & bit) != 0 { + index += 1; + } + index +} + + #[cfg(test)] +mod test { + + use image::ImageError; + + use super::*; + + #[test] + fn test_big_image() -> Result<(), ImageError> { + let src: RgbaImage = image::open("testdata/input.jpg").unwrap().into_rgba8(); + + let out = quantize(&src, 256); + out.save_with_format("output.jpg", image::ImageFormat::Jpeg)?; + Ok(()) + } +} diff --git a/testdata/img.png b/testdata/img.png new file mode 100644 index 0000000..bcb1463 Binary files /dev/null and b/testdata/img.png differ diff --git a/testdata/input.jpg b/testdata/input.jpg new file mode 100644 index 0000000..c8417a6 Binary files /dev/null and b/testdata/input.jpg differ diff --git a/testdata/test.jpg b/testdata/test.jpg new file mode 100644 index 0000000..7f758f6 Binary files /dev/null and b/testdata/test.jpg differ