New node: Dehaze (#1882)

* feat: Implemented Dehaze Node * Update Cargo.toml * Remove unecessary image conversions * Code review * Further fixes --------- Co-authored-by: Dennis Kobert <dennis@kobert.dev> Co-authored-by: Keavon Chambers <keavon@keavon.com>
2024-09-25 07:44:48 +04:00 · 2024-09-25 07:44:48 +04:00 · c5454af48b
parent 0b0169a415
commit c5454af48b
4 changed files with 345 additions and 0 deletions
--- a/Cargo.lock
+++ b/Cargo.lock
@ -2480,6 +2480,7 @@ dependencies = [
 "image-compare",
 "js-sys",
 "log",
+ "ndarray",
 "node-macro",
 "path-bool",
 "rand 0.8.5",
@ -3605,6 +3606,16 @@ version = "0.7.3"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "0e7465ac9959cc2b1404e8e2367b43684a6d13790fe23056cc8c6c5a6b7bcb94"

+[[package]]
+name = "matrixmultiply"
+version = "0.3.8"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "7574c1cf36da4798ab73da5b215bbf444f50718207754cb522201d78d1cd0ff2"
+dependencies = [
+ "autocfg",
+ "rawpointer",
+]
+
 [[package]]
 name = "memchr"
 version = "2.7.4"
@ -3744,6 +3755,21 @@ dependencies = [
 "tempfile",
 ]

+[[package]]
+name = "ndarray"
+version = "0.16.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "882ed72dce9365842bf196bdeedf5055305f11fc8c03dee7bb0194a6cad34841"
+dependencies = [
+ "matrixmultiply",
+ "num-complex",
+ "num-integer",
+ "num-traits",
+ "portable-atomic",
+ "portable-atomic-util",
+ "rawpointer",
+]
+
 [[package]]
 name = "ndk"
 version = "0.6.0"
@ -3892,6 +3918,15 @@ dependencies = [
 "winapi",
 ]

+[[package]]
+name = "num-complex"
+version = "0.4.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "73f88a1307638156682bada9d7604135552957b7818057dcef22705b4d509495"
+dependencies = [
+ "num-traits",
+]
+
 [[package]]
 name = "num-conv"
 version = "0.1.0"
@ -3909,6 +3944,15 @@ dependencies = [
 "syn 2.0.77",
 ]

+[[package]]
+name = "num-integer"
+version = "0.1.46"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "7969661fd2958a5cb096e56c8e1ad0444ac2bbcd0061bd28660485a44879858f"
+dependencies = [
+ "num-traits",
+]
+
 [[package]]
 name = "num-traits"
 version = "0.2.19"
@ -4618,6 +4662,21 @@ dependencies = [
 "windows-sys 0.59.0",
 ]

+[[package]]
+name = "portable-atomic"
+version = "1.8.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "d30538d42559de6b034bc76fd6dd4c38961b1ee5c6c56e3808c50128fdbc22ce"
+
+[[package]]
+name = "portable-atomic-util"
+version = "0.2.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "fcdd8420072e66d54a407b3316991fe946ce3ab1083a7f575b2463866624704d"
+dependencies = [
+ "portable-atomic",
+]
+
 [[package]]
 name = "powerfmt"
 version = "0.2.0"
@ -4916,6 +4975,12 @@ version = "0.6.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "20675572f6f24e9e76ef639bc5552774ed45f1c30e2951e1e99c59888861c539"

+[[package]]
+name = "rawpointer"
+version = "0.2.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "60a357793950651c4ed0f3f52338f53b2f809f32d83a07f72909fa13e4c6c1e3"
+
 [[package]]
 name = "rayon"
 version = "1.10.0"
--- a/node-graph/gstd/Cargo.toml
+++ b/node-graph/gstd/Cargo.toml
@ -88,3 +88,4 @@ web-sys = { workspace = true, optional = true, features = [

 # Optional dependencies
 image-compare = { version = "0.4.1", optional = true }
+ndarray = "0.16.1"
--- a/node-graph/gstd/src/dehaze.rs
+++ b/node-graph/gstd/src/dehaze.rs
@ -0,0 +1,277 @@
+use graph_craft::proto::types::Percentage;
+use graphene_core::raster::{Image, ImageFrame};
+use graphene_core::transform::Footprint;
+use graphene_core::Color;
+
+use image::{DynamicImage, GenericImage, GenericImageView, GrayImage, ImageBuffer, Luma, Rgba, RgbaImage};
+use ndarray::{Array2, ArrayBase, Dim, OwnedRepr};
+use std::cmp::{max, min};
+
+#[node_macro::node(category("Raster: Filter"))]
+async fn dehaze<F: 'n + Send + Sync>(
+	#[implementations(
+		(),
+		Footprint,
+	)]
+	footprint: F,
+	#[implementations(
+		() -> ImageFrame<Color>,
+		Footprint -> ImageFrame<Color>,
+	)]
+	image_frame: impl Node<F, Output = ImageFrame<Color>>,
+	strength: Percentage,
+) -> ImageFrame<Color> {
+	let image_frame = image_frame.eval(footprint).await;
+
+	// Prepare the image data for processing
+	let image = image_frame.image;
+	let image_data = bytemuck::cast_vec(image.data);
+	let image_buffer = image::Rgba32FImage::from_raw(image.width, image.height, image_data).expect("Failed to convert internal ImageFrame into image-rs data type.");
+	let dynamic_image: image::DynamicImage = image_buffer.into();
+
+	// Run the dehaze algorithm
+	let dehazed_dynamic_image = dehaze_image(dynamic_image, strength / 100.);
+
+	// Prepare the image data for returning
+	let buffer = dehazed_dynamic_image.to_rgba32f().into_raw();
+	let color_vec = bytemuck::cast_vec(buffer);
+	let dehazed_image = Image {
+		width: image.width,
+		height: image.height,
+		data: color_vec,
+		base64_string: None,
+	};
+
+	ImageFrame {
+		image: dehazed_image,
+		transform: image_frame.transform,
+		alpha_blending: image_frame.alpha_blending,
+	}
+}
+
+// There is no real point in modifying these values because they do not change the final result all that much.
+// The authors of the paper recommended using these values to get a reasonable balance of performance and quality.
+const PATCH_SIZE: u32 = 15;
+const TOP_PERCENT: f64 = 0.001;
+const RADIUS: u32 = 60;
+const EPSILON: f64 = 0.0001;
+const TX: f32 = 0.1;
+
+// Dehazing algorithm based on "Single Image Haze Removal Using Dark Channel Prior"
+// Paper: <https://www.researchgate.net/publication/220182411_Single_Image_Haze_Removal_Using_Dark_Channel_Prior>
+// TODO: Make this algorithm work with negative strength values
+fn dehaze_image(image: DynamicImage, strength: f64) -> DynamicImage {
+	// TODO: Break out this pair of steps into its own node, with a memoize node which caches the pair of outputs, so the strength can be adjusted without recomputing these two steps.
+	let dark_channel = compute_dark_channel(&image);
+	let atmospheric_light = estimate_atmospheric_light(&image, &dark_channel);
+
+	let transmission_map = estimate_transmission_map(&image, &dark_channel, strength);
+	let refined_transmission_map = refine_transmission_map(&image, &transmission_map);
+
+	recover(&image, &refined_transmission_map, atmospheric_light)
+}
+
+fn compute_dark_channel(image: &DynamicImage) -> DynamicImage {
+	let (width, height) = image.dimensions();
+	let mut dark_channel = GrayImage::new(width, height);
+	let half_patch = PATCH_SIZE / 2;
+
+	for y in 0..height {
+		for x in 0..width {
+			let pixel = image.get_pixel(x, y);
+			let min_intensity = min(min(pixel[0], pixel[1]), pixel[2]);
+			dark_channel.put_pixel(x, y, Luma([min_intensity]));
+		}
+	}
+
+	let mut eroded_channel = RgbaImage::new(width, height);
+
+	for y in 0..height {
+		for x in 0..width {
+			let mut local_min = u8::MAX;
+
+			for dy in 0..PATCH_SIZE {
+				for dx in 0..PATCH_SIZE {
+					let nx = x as i32 + dx as i32 - half_patch as i32;
+					let ny = y as i32 + dy as i32 - half_patch as i32;
+
+					if nx >= 0 && nx < width as i32 && ny >= 0 && ny < height as i32 {
+						let intensity = dark_channel.get_pixel(nx as u32, ny as u32)[0];
+						if intensity < local_min {
+							local_min = intensity;
+						}
+					}
+				}
+			}
+			let alpha = image.get_pixel(x, y)[3];
+			eroded_channel.put_pixel(x, y, Rgba([local_min, local_min, local_min, alpha]));
+		}
+	}
+
+	DynamicImage::ImageRgba8(eroded_channel)
+}
+
+fn estimate_atmospheric_light(hazy: &DynamicImage, dark_channel: &DynamicImage) -> Rgba<u8> {
+	let (width, height) = hazy.dimensions();
+	let dark = dark_channel.to_luma_alpha8();
+	let total_pixels = (width * height) as usize;
+	let num_pixels = ((TOP_PERCENT / 100.) * total_pixels as f64).ceil() as usize;
+
+	let mut intensities: Vec<(u32, u32, f64)> = Vec::with_capacity(total_pixels);
+
+	for y in 0..height {
+		for x in 0..width {
+			let pixel = dark.get_pixel(x, y);
+			let intensity = pixel.0[0] as f64;
+			intensities.push((x, y, intensity))
+		}
+	}
+
+	intensities.sort_by(|a, b| b.2.partial_cmp(&a.2).unwrap());
+
+	let top_intensities = &intensities[..num_pixels];
+
+	let mut atm_sum = [0., 0., 0.];
+	for (x, y, _) in top_intensities {
+		let pixel = hazy.get_pixel(*x, *y);
+		atm_sum[0] += pixel[0] as f64;
+		atm_sum[1] += pixel[1] as f64;
+		atm_sum[2] += pixel[2] as f64;
+	}
+
+	let num_pixels = num_pixels as f64;
+
+	Rgba([(atm_sum[0] / num_pixels) as u8, (atm_sum[1] / num_pixels) as u8, (atm_sum[2] / num_pixels) as u8, 255])
+}
+
+fn estimate_transmission_map(image: &DynamicImage, dark_channel: &DynamicImage, omega: f64) -> DynamicImage {
+	let (width, height) = image.dimensions();
+	let mut transmission_map = RgbaImage::new(width, height);
+
+	for y in 0..height {
+		for x in 0..width {
+			let min_intensity = dark_channel.get_pixel(x, y).0[0] as f32 / 255.;
+			let transmission_value = 1. - omega * min_intensity as f64;
+			let alpha = image.get_pixel(x, y)[3];
+			transmission_map.put_pixel(
+				x,
+				y,
+				Rgba([(transmission_value * 255.) as u8, (transmission_value * 255.) as u8, (transmission_value * 255.) as u8, alpha]),
+			);
+		}
+	}
+
+	DynamicImage::ImageRgba8(transmission_map)
+}
+
+fn refine_transmission_map(img: &DynamicImage, transmission_map: &DynamicImage) -> DynamicImage {
+	let gray_image = img.to_luma8();
+
+	let normalized_gray_image: GrayImage = ImageBuffer::from_fn(gray_image.width(), gray_image.height(), |x, y| {
+		let pixel = gray_image.get_pixel(x, y);
+		let normalized_value = (pixel[0] as f64 / 255.) * 255.;
+		Luma([normalized_value as u8])
+	});
+
+	let normalized_gray_image = DynamicImage::ImageLuma8(normalized_gray_image);
+
+	guided_filter(&normalized_gray_image, transmission_map, RADIUS, EPSILON)
+}
+
+fn recover(im: &DynamicImage, t: &DynamicImage, a: Rgba<u8>) -> DynamicImage {
+	let (width, height) = im.dimensions();
+	let mut res = DynamicImage::new_rgba8(width, height);
+
+	let a = [a[0] as f32 / 255., a[1] as f32 / 255., a[2] as f32 / 255.];
+
+	for y in 0..height {
+		for x in 0..width {
+			let im_pixel = im.get_pixel(x, y).0;
+			let t_pixel = t.get_pixel(x, y).0;
+			let t_val = f32::max(t_pixel[0] as f32 / 255., TX);
+
+			let mut res_pixel = [0; 4];
+			for ind in 0..3 {
+				res_pixel[ind] = ((((im_pixel[ind] as f32 / 255. - a[ind]) / t_val) + a[ind]).clamp(0., 1.) * 255.) as u8;
+			}
+			res_pixel[3] = im_pixel[3];
+
+			res.put_pixel(x, y, Rgba(res_pixel));
+		}
+	}
+
+	res
+}
+
+fn guided_filter(guidance_img: &DynamicImage, input_img: &DynamicImage, r: u32, epsilon: f64) -> DynamicImage {
+	let (width, height) = guidance_img.dimensions();
+	let radius = r as i32;
+
+	let guidance_nd = image_to_ndarray(guidance_img);
+	let input_nd = image_to_ndarray(input_img);
+
+	let mean_guidance = box_filter(&guidance_nd, radius);
+	let mean_input = box_filter(&input_nd, radius);
+	let corr_guidance = box_filter(&(guidance_nd.clone() * guidance_nd.clone()), radius);
+	let corr_guidance_input = box_filter(&(guidance_nd.clone() * input_nd.clone()), radius);
+
+	let var_guidance = &corr_guidance - &(mean_guidance.clone() * mean_guidance.clone());
+	let cov_guidance_input = &corr_guidance_input - &(mean_guidance.clone() * mean_input.clone());
+
+	let a = &cov_guidance_input / &(var_guidance.clone() + epsilon);
+	let b = mean_input - &(a.clone() * mean_guidance);
+
+	let mean_a = box_filter(&a, radius);
+	let mean_b = box_filter(&b, radius);
+
+	let q = &mean_a * &guidance_nd + mean_b;
+
+	ndarray_to_image(&q, width, height)
+}
+
+fn box_filter(img: &Array2<f64>, radius: i32) -> Array2<f64> {
+	let (height, width) = img.dim();
+	let mut result = Array2::zeros((height, width));
+	let mut integral_image: ArrayBase<OwnedRepr<f64>, Dim<[usize; 2]>> = Array2::zeros((height + 1, width + 1));
+
+	// Compute integral image
+	for y in 0..height {
+		for x in 0..width {
+			integral_image[(y + 1, x + 1)] = img[(y, x)] + integral_image[(y, x + 1)] + integral_image[(y + 1, x)] - integral_image[(y, x)];
+		}
+	}
+
+	for y in 0..height {
+		for x in 0..width {
+			let y1 = max(0, y as i32 - radius) as usize;
+			let y2 = min(height as i32 - 1, y as i32 + radius) as usize;
+			let x1 = max(0, x as i32 - radius) as usize;
+			let x2 = min(width as i32 - 1, x as i32 + radius) as usize;
+
+			let area = (y2 - y1 + 1) as f64 * (x2 - x1 + 1) as f64;
+
+			result[(y, x)] = (integral_image[(y2 + 1, x2 + 1)] - integral_image[(y1, x2 + 1)] - integral_image[(y2 + 1, x1)] + integral_image[(y1, x1)]) / area;
+		}
+	}
+
+	result
+}
+
+fn image_to_ndarray(img: &DynamicImage) -> Array2<f64> {
+	let (width, height) = img.dimensions();
+	let mut array = Array2::zeros((height as usize, width as usize));
+	for (x, y, pixel) in img.pixels() {
+		let luminance = pixel.0[0] as f64 / 255.;
+		array[(y as usize, x as usize)] = luminance;
+	}
+	array
+}
+
+fn ndarray_to_image(array: &Array2<f64>, width: u32, height: u32) -> DynamicImage {
+	let mut img = DynamicImage::new_rgba8(width, height);
+	for ((y, x), &value) in array.indexed_iter() {
+		let clamped_value = (value * 255.).clamp(0., 255.) as u8;
+		img.put_pixel(x as u32, y as u32, Rgba([clamped_value, clamped_value, clamped_value, 255]));
+	}
+	img
+}
--- a/node-graph/gstd/src/lib.rs
+++ b/node-graph/gstd/src/lib.rs
@ -27,4 +27,6 @@ pub mod brush;
 #[cfg(feature = "wasm")]
 pub mod wasm_application_io;

+pub mod dehaze;
+
 pub mod imaginate;