Quantization node (#887)

Add basic quantization node

node-graph/gcore/src/raster.rs

@@ -3,7 +3,7 @@ use crate::Node;
 pub mod color;
 pub use self::color::Color;
 
-#[derive(Debug, Clone, Copy)]
+#[derive(Debug, Clone, Copy, Default)]
 pub struct GrayscaleColorNode;
 
 impl Node<Color> for GrayscaleColorNode {
@@ -53,6 +53,32 @@ impl<N: Node<(), Output = f32> + Copy> BrightenColorNode<N> {
 	}
 }
 
+#[derive(Debug, Clone, Copy)]
+pub struct GammaColorNode<N: Node<(), Output = f32>>(N);
+
+impl<N: Node<(), Output = f32>> Node<Color> for GammaColorNode<N> {
+	type Output = Color;
+	fn eval(self, color: Color) -> Color {
+		let gamma = self.0.eval(());
+		let per_channel = |col: f32| col.powf(gamma);
+		Color::from_rgbaf32_unchecked(per_channel(color.r()), per_channel(color.g()), per_channel(color.b()), color.a())
+	}
+}
+impl<N: Node<(), Output = f32> + Copy> Node<Color> for &GammaColorNode<N> {
+	type Output = Color;
+	fn eval(self, color: Color) -> Color {
+		let gamma = self.0.eval(());
+		let per_channel = |col: f32| col.powf(gamma);
+		Color::from_rgbaf32_unchecked(per_channel(color.r()), per_channel(color.g()), per_channel(color.b()), color.a())
+	}
+}
+
+impl<N: Node<(), Output = f32> + Copy> GammaColorNode<N> {
+	pub fn new(node: N) -> Self {
+		Self(node)
+	}
+}
+
 #[derive(Debug, Clone, Copy)]
 #[cfg(not(target_arch = "spirv"))]
 pub struct HueShiftColorNode<N: Node<(), Output = f32>>(N);

node-graph/gstd/Cargo.toml

@@ -14,9 +14,11 @@ default = ["memoization"]
 gpu = ["graphene-core/gpu", "gpu-compiler-bin-wrapper", "compilation-client"]
 vulkan = ["gpu", "vulkan-executor"]
 wgpu = ["gpu", "wgpu-executor"]
+quantization = ["autoquant"]
 
 
 [dependencies]
+autoquant = { git = "https://github.com/truedoctor/autoquant", optional = true, features = ["fitting"] }
 graphene-core = {path = "../gcore", features = ["async", "std" ], default-features = false}
 borrow_stack = {path = "../borrow_stack"}
 dyn-any = {path = "../../libraries/dyn-any", features = ["derive"]}

node-graph/gstd/src/lib.rs

@@ -15,4 +15,7 @@ pub mod any;
 #[cfg(feature = "gpu")]
 pub mod executor;
 
+#[cfg(feature = "quantization")]
+pub mod quantization;
+
 pub use graphene_core::*;

node-graph/gstd/src/quantization.rs

@@ -0,0 +1,54 @@
+use graphene_core::raster::{Color, Image};
+use graphene_core::Node;
+
+/// The `GenerateQuantizationNode` encodes the brightness of each channel of the image as an integer number
+/// sepified by the samples parameter. This node is used to asses the loss of visual information when
+/// quantizing the image using different fit functions.
+pub struct GenerateQuantizationNode<N: Node<(), Output = u32>, M: Node<(), Output = u32>> {
+	samples: N,
+	function: M,
+}
+
+#[node_macro::node_fn(GenerateQuantizationNode)]
+fn generate_quantization_fn(image: Image, samples: u32, function: u32) -> Image {
+	// Scale the input image, this can be removed by adding an extra parameter to the fit function.
+	let max_energy = 16380.;
+	let data: Vec<f64> = image.data.iter().flat_map(|x| vec![x.r() as f64, x.g() as f64, x.b() as f64]).collect();
+	let data: Vec<f64> = data.iter().map(|x| x * max_energy).collect();
+	let mut dist = autoquant::integrate_distribution(data);
+	autoquant::drop_duplicates(&mut dist);
+	let dist = autoquant::normalize_distribution(dist.as_slice());
+	let max = dist.last().unwrap().0;
+	let linear = Box::new(autoquant::SimpleFitFn {
+		function: move |x| x / max,
+		inverse: move |x| x * max,
+		name: "identity",
+	});
+	let best = match function {
+		0 => linear as Box<dyn autoquant::FitFn>,
+		1 => linear as Box<dyn autoquant::FitFn>,
+		2 => Box::new(autoquant::models::OptimizedLog::new(dist, 20)) as Box<dyn autoquant::FitFn>,
+		_ => linear as Box<dyn autoquant::FitFn>,
+	};
+
+	let roundtrip = |sample: f32| -> f32 {
+		let encoded = autoquant::encode(sample as f64 * max_energy, best.as_ref(), samples);
+		let decoded = autoquant::decode(encoded, best.as_ref(), samples) / max_energy;
+		log::trace!("{} enc: {} dec: {}", sample, encoded, decoded);
+		decoded as f32
+	};
+
+	let new_data = image
+		.data
+		.iter()
+		.map(|c| {
+			let r = roundtrip(c.r());
+			let g = roundtrip(c.g());
+			let b = roundtrip(c.b());
+			let a = c.a();
+
+			Color::from_rgbaf32_unchecked(r, g, b, a)
+		})
+		.collect();
+	Image { data: new_data, ..image }
+}

node-graph/interpreted-executor/Cargo.toml

@@ -8,6 +8,7 @@ license = "MIT OR Apache-2.0"
 default = []
 serde = ["dep:serde", "graphene-std/serde", "glam/serde"]
 gpu = ["graphene-std/gpu", "graphene-core/gpu", "graphene-std/wgpu"]
+quantization = ["graphene-std/quantization"]
 
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
 

node-graph/interpreted-executor/src/node_registry.rs

@@ -268,6 +268,32 @@ static NODE_REGISTRY: &[(NodeIdentifier, NodeConstructor)] = &[
 			}
 		},
 	),
+	#[cfg(feature = "quantization")]
+	(
+		NodeIdentifier::new("graphene_std::quantization::GenerateQuantizationNode", &[concrete!("&TypeErasedNode")]),
+		|proto_node, stack| {
+			if let ConstructionArgs::Nodes(operation_node_id) = proto_node.construction_args {
+				stack.push_fn(move |nodes| {
+					info!("Quantization Depending upon id {:?}", operation_node_id);
+					let samples_node = nodes.get(operation_node_id[0] as usize).unwrap();
+					let index_node = nodes.get(operation_node_id[1] as usize).unwrap();
+					let samples_node: DowncastBothNode<_, (), u32> = DowncastBothNode::new(samples_node);
+					let index_node: DowncastBothNode<_, (), u32> = DowncastBothNode::new(index_node);
+					let map_node = graphene_std::quantization::GenerateQuantizationNode::new(samples_node, index_node);
+					let map_node = DynAnyNode::new(map_node);
+
+					if let ProtoNodeInput::Node(node_id) = proto_node.input {
+						let pre_node = nodes.get(node_id as usize).unwrap();
+						(pre_node).then(map_node).into_type_erased()
+					} else {
+						map_node.into_type_erased()
+					}
+				})
+			} else {
+				unimplemented!()
+			}
+		},
+	),
 	(NodeIdentifier::new("graphene_std::raster::MapImageNode", &[]), |proto_node, stack| {
 		if let ConstructionArgs::Nodes(operation_node_id) = proto_node.construction_args {
 			stack.push_fn(move |nodes| {