Add the image segmentation node (#1122)

* Add image segmentation node

This is a node which receives as input:
- An image.
- A mask, which consists of colors that define the image segments.

Each unique color in the mask defines an area where a segmen resides.
The node generates a `Vec<ImageFrame>` where the length of the result
is the number of unique colors in the mask.

Signed-off-by: Ică Alexandru-Gabriel <alexandru@seyhanlee.com>

* Add the Index node for image segments

Since the output of the image segmentation node is a `Vec<ImageFrame>`,
we want a way to access the segments individually. The Index node receives
a `Vec<ImageFrame>` as input and an index, and returns the image found
at that index in the vec.

Signed-off-by: Ică Alexandru-Gabriel <alexandru@seyhanlee.com>

* Integrate the image segmentation and index nodes into the editor

Signed-off-by: Ică Alexandru-Gabriel <alexandru@seyhanlee.com>

* Initialize the input of the index node with an empty image frame

Signed-off-by: Ică Alexandru-Gabriel <alexandru@seyhanlee.com>

* Don't expose the parameter for the index node

Signed-off-by: Ică Alexandru-Gabriel <alexandru@seyhanlee.com>

* Don't crash the editor when the number of segments exceeds the accepted limit

Signed-off-by: Ică Alexandru-Gabriel <alexandru@seyhanlee.com>

* Print a warning in the console when the number of segments exceeds the accepted limit

Signed-off-by: Ică Alexandru-Gabriel <alexandru@seyhanlee.com>

* Add a few more checks so that the editor doesn't crash on invalid input

Signed-off-by: Ică Alexandru-Gabriel <alexandru@seyhanlee.com>

* Replace the tagged value for the index node

Signed-off-by: Ică Alexandru-Gabriel <alexandru@seyhanlee.com>

* Fix merge conflicts

---------

Signed-off-by: Ică Alexandru-Gabriel <alexandru@seyhanlee.com>
Co-authored-by: Dennis Kobert <dennis@kobert.dev>

node-graph/gstd/src/image_segmentation.rs

@@ -0,0 +1,126 @@
+use std::collections::hash_map::HashMap;
+
+use graphene_core::raster::{Color, ImageFrame};
+use graphene_core::Node;
+
+fn apply_mask(image_frame: &mut ImageFrame<Color>, x: usize, y: usize, multiplier: u8) {
+	let color = &mut image_frame.image.data[y * image_frame.image.width as usize + x];
+	let color8 = color.to_rgba8();
+	*color = Color::from_rgba8_srgb(color8[0] * multiplier, color8[1] * multiplier, color8[2] * multiplier, color8[3] * multiplier);
+}
+
+pub struct Mask {
+	pub data: Vec<u8>,
+	pub width: usize,
+	pub height: usize,
+}
+
+impl Mask {
+	fn sample(&self, u: f32, v: f32) -> u8 {
+		let x = (u * (self.width as f32)) as usize;
+		let y = (v * (self.height as f32)) as usize;
+
+		self.data[y * self.width + x]
+	}
+}
+
+fn image_segmentation(input_image: &ImageFrame<Color>, input_mask: &Mask) -> Vec<ImageFrame<Color>> {
+	const NUM_LABELS: usize = u8::MAX as usize;
+	let mut result = Vec::<ImageFrame<Color>>::with_capacity(NUM_LABELS);
+	let mut current_label = 0_usize;
+	let mut label_appeared = [false; NUM_LABELS + 1];
+	let mut max_label = 0_usize;
+
+	if input_mask.data.is_empty() {
+		warn!("The mask for the segmentation node is empty!");
+		return vec![ImageFrame::empty()];
+	}
+
+	result.push(input_image.clone());
+	let result_last = result.last_mut().unwrap();
+
+	for y in 0..input_image.image.height {
+		let v = (y as f32) / (input_image.image.height as f32);
+		for x in 0..input_image.image.width {
+			let u = (x as f32) / (input_image.image.width as f32);
+			let label = input_mask.sample(u, v) as usize;
+			let multiplier = (label == current_label) as u8;
+
+			apply_mask(result_last, x as usize, y as usize, multiplier);
+
+			if label < NUM_LABELS {
+				label_appeared[label] = true;
+				max_label = max_label.max(label);
+			}
+		}
+	}
+
+	if !label_appeared[current_label] {
+		result.pop();
+	}
+
+	for i in 1..=max_label.max(NUM_LABELS) {
+		current_label = i;
+
+		if !label_appeared[current_label] {
+			continue;
+		}
+
+		result.push(input_image.clone());
+		let result_last = result.last_mut().unwrap();
+
+		for y in 0..input_image.image.height {
+			let v = (y as f32) / (input_image.image.height as f32);
+			for x in 0..input_image.image.width {
+				let u = (x as f32) / (input_image.image.width as f32);
+				let label = input_mask.sample(u, v) as usize;
+				let multiplier = (label == current_label) as u8;
+
+				apply_mask(result_last, x as usize, y as usize, multiplier);
+			}
+		}
+	}
+
+	result
+}
+
+fn convert_image_to_mask(input: &ImageFrame<Color>) -> Vec<u8> {
+	let mut result = vec![0_u8; (input.image.width * input.image.height) as usize];
+	let mut colors = HashMap::<[u8; 4], usize>::new();
+	let mut last_value = 0_usize;
+
+	for (color, result) in input.image.data.iter().zip(result.iter_mut()) {
+		let color = color.to_rgba8();
+		if let Some(value) = colors.get(&color) {
+			*result = *value as u8;
+		} else {
+			if last_value > u8::MAX as usize {
+				warn!("The limit for number of segments ({}) has been exceeded!", u8::MAX);
+				break;
+			}
+
+			*result = last_value as u8;
+			colors.insert(color, last_value);
+			last_value += 1;
+		}
+	}
+
+	result
+}
+
+#[derive(Debug)]
+pub struct ImageSegmentationNode<MaskImage> {
+	pub(crate) mask_image: MaskImage,
+}
+
+#[node_macro::node_fn(ImageSegmentationNode)]
+pub(crate) fn image_segmentation(image: ImageFrame<Color>, mask_image: ImageFrame<Color>) -> Vec<ImageFrame<Color>> {
+	let mask_data = convert_image_to_mask(&mask_image);
+	let mask = Mask {
+		data: mask_data,
+		width: mask_image.image.width as usize,
+		height: mask_image.image.height as usize,
+	};
+
+	image_segmentation(&image, &mask)
+}

node-graph/gstd/src/lib.rs

@@ -20,4 +20,6 @@ pub mod quantization;
 
 pub use graphene_core::*;
 
+pub mod image_segmentation;
+
 pub mod brush;