Raw-rs: Refactor to run multiple steps in a single loop (#1972)

* Prevent extra allocation in convert to RGB step

* Run preprocessing steps in a single loop

* Create new API to call steps in pipeline

* Include transform and gamma correction step

* cargo fmt

* Split scale colors into two steps

* Code relocations

* cargo fmt

* Implement transform traits for all tuples

* Replace Captures trick with the new `use` keyword

libraries/raw-rs/src/preprocessing/camera_data.rs

@@ -1,119 +0,0 @@
-use crate::RawImage;
-use build_camera_data::build_camera_data;
-
-pub struct CameraData {
-	pub black: u16,
-	pub maximum: u16,
-	pub camera_to_xyz: [i16; 9],
-}
-
-impl CameraData {
-	const DEFAULT: CameraData = CameraData {
-		black: 0,
-		maximum: 0,
-		camera_to_xyz: [0; 9],
-	};
-}
-
-const CAMERA_DATA: [(&str, CameraData); 40] = build_camera_data!();
-
-const XYZ_TO_RGB: [[f64; 3]; 3] = [
-	// Matrix:
-	[0.412453, 0.357580, 0.180423],
-	[0.212671, 0.715160, 0.072169],
-	[0.019334, 0.119193, 0.950227],
-];
-
-pub fn calculate_conversion_matrices(mut raw_image: RawImage) -> RawImage {
-	let Some(ref camera_model) = raw_image.camera_model else { return raw_image };
-	let camera_name_needle = camera_model.make.to_owned() + " " + &camera_model.model;
-
-	let camera_to_xyz = CAMERA_DATA
-		.iter()
-		.find(|(camera_name_haystack, _)| camera_name_needle == *camera_name_haystack)
-		.map(|(_, data)| data.camera_to_xyz.map(|x| (x as f64) / 10_000.));
-	let Some(camera_to_xyz) = camera_to_xyz else { return raw_image };
-
-	let mut camera_to_rgb = [[0.; 3]; 3];
-	for i in 0..3 {
-		for j in 0..3 {
-			for k in 0..3 {
-				camera_to_rgb[i][j] += camera_to_xyz[i * 3 + k] * XYZ_TO_RGB[k][j];
-			}
-		}
-	}
-
-	let white_balance_multiplier = camera_to_rgb.map(|x| 1. / x.iter().sum::<f64>());
-	for (index, row) in camera_to_rgb.iter_mut().enumerate() {
-		*row = row.map(|x| x * white_balance_multiplier[index]);
-	}
-	let rgb_to_camera = transpose(pseudoinverse(camera_to_rgb));
-
-	let cfa_white_balance_multiplier = if let Some(white_balance) = raw_image.camera_white_balance_multiplier {
-		white_balance
-	} else {
-		raw_image.cfa_pattern.map(|index| white_balance_multiplier[index as usize])
-	};
-
-	raw_image.white_balance_multiplier = Some(cfa_white_balance_multiplier);
-	raw_image.camera_to_rgb = Some(camera_to_rgb);
-	raw_image.rgb_to_camera = Some(rgb_to_camera);
-
-	raw_image
-}
-
-#[allow(clippy::needless_range_loop)]
-fn pseudoinverse<const N: usize>(matrix: [[f64; 3]; N]) -> [[f64; 3]; N] {
-	let mut output_matrix = [[0.; 3]; N];
-	let mut work = [[0.; 6]; 3];
-
-	for i in 0..3 {
-		for j in 0..6 {
-			work[i][j] = if j == i + 3 { 1. } else { 0. };
-		}
-		for j in 0..3 {
-			for k in 0..N {
-				work[i][j] += matrix[k][i] * matrix[k][j];
-			}
-		}
-	}
-
-	for i in 0..3 {
-		let num = work[i][i];
-		for j in 0..6 {
-			work[i][j] /= num;
-		}
-		for k in 0..3 {
-			if k == i {
-				continue;
-			}
-			let num = work[k][i];
-			for j in 0..6 {
-				work[k][j] -= work[i][j] * num;
-			}
-		}
-	}
-
-	for i in 0..N {
-		for j in 0..3 {
-			output_matrix[i][j] = 0.;
-			for k in 0..3 {
-				output_matrix[i][j] += work[j][k + 3] * matrix[i][k];
-			}
-		}
-	}
-
-	output_matrix
-}
-
-fn transpose<const N: usize>(matrix: [[f64; 3]; N]) -> [[f64; N]; 3] {
-	let mut output_matrix = [[0.; N]; 3];
-
-	for (i, row) in matrix.iter().enumerate() {
-		for (j, &value) in row.iter().enumerate() {
-			output_matrix[j][i] = value;
-		}
-	}
-
-	output_matrix
-}

libraries/raw-rs/src/preprocessing/mod.rs

@@ -1,3 +1,3 @@
-pub mod camera_data;
-pub mod scale_colors;
+pub mod scale_to_16bit;
+pub mod scale_white_balance;
 pub mod subtract_black;

libraries/raw-rs/src/preprocessing/scale_colors.rs

@@ -1,37 +0,0 @@
-use crate::RawImage;
-
-pub fn scale_colors(mut raw_image: RawImage) -> RawImage {
-	let Some(mut white_balance_multiplier) = raw_image.white_balance_multiplier else {
-		return raw_image;
-	};
-
-	if white_balance_multiplier[1] == 0. {
-		white_balance_multiplier[1] = 1.;
-	}
-
-	// TODO: Move this at its correct location when highlights are implemented correctly.
-	let highlight = 0;
-
-	let normalize_white_balance = if highlight == 0 {
-		white_balance_multiplier.iter().copied().fold(f64::INFINITY, f64::min)
-	} else {
-		white_balance_multiplier.iter().copied().fold(f64::NEG_INFINITY, f64::max)
-	};
-
-	let final_multiplier = if normalize_white_balance > 0.00001 && raw_image.maximum > 0 {
-		let scale_to_16bit_multiplier = u16::MAX as f64 / raw_image.maximum as f64;
-		white_balance_multiplier.map(|x| x / normalize_white_balance * scale_to_16bit_multiplier)
-	} else {
-		[1., 1., 1., 1.]
-	};
-
-	for row in 0..raw_image.height {
-		for column in 0..raw_image.width {
-			let index = row * raw_image.width + column;
-			let cfa_index = 2 * (row % 2) + (column % 2);
-			raw_image.data[index] = ((raw_image.data[index] as f64) * final_multiplier[cfa_index]).min(u16::MAX as f64).max(0.) as u16;
-		}
-	}
-
-	raw_image
-}

libraries/raw-rs/src/preprocessing/scale_to_16bit.rs

@@ -0,0 +1,15 @@
+use crate::{RawImage, RawPixel, SubtractBlack};
+
+impl RawImage {
+	pub fn scale_to_16bit_fn(&self) -> impl Fn(RawPixel) -> u16 {
+		let black_level = match self.black {
+			SubtractBlack::CfaGrid(x) => x,
+			_ => unreachable!(),
+		};
+
+		let maximum = self.maximum - black_level.iter().max().unwrap();
+		let scale_to_16bit_multiplier = if maximum > 0 { u16::MAX as f64 / maximum as f64 } else { 1. };
+
+		move |pixel: RawPixel| ((pixel.value as f64) * scale_to_16bit_multiplier).min(u16::MAX as f64).max(0.) as u16
+	}
+}

libraries/raw-rs/src/preprocessing/scale_white_balance.rs

@@ -0,0 +1,31 @@
+use crate::{RawImage, RawPixel};
+
+impl RawImage {
+	pub fn scale_white_balance_fn(&self) -> impl Fn(RawPixel) -> u16 {
+		let Some(mut white_balance) = self.white_balance else { todo!() };
+
+		if white_balance[1] == 0. {
+			white_balance[1] = 1.;
+		}
+
+		// TODO: Move this at its correct location when highlights are implemented correctly.
+		let highlight = 0;
+
+		let normalization_factor = if highlight == 0 {
+			white_balance.into_iter().fold(f64::INFINITY, f64::min)
+		} else {
+			white_balance.into_iter().fold(f64::NEG_INFINITY, f64::max)
+		};
+
+		let normalized_white_balance = if normalization_factor > 0.00001 {
+			white_balance.map(|x| x / normalization_factor)
+		} else {
+			[1., 1., 1., 1.]
+		};
+
+		move |pixel: RawPixel| {
+			let cfa_index = 2 * (pixel.row % 2) + (pixel.column % 2);
+			((pixel.value as f64) * normalized_white_balance[cfa_index]).min(u16::MAX as f64).max(0.) as u16
+		}
+	}
+}

libraries/raw-rs/src/preprocessing/subtract_black.rs

@@ -1,30 +1,11 @@
+use crate::RawPixel;
 use crate::{RawImage, SubtractBlack};
 
-pub fn subtract_black(raw_image: RawImage) -> RawImage {
-	let mut raw_image = match raw_image.black {
-		SubtractBlack::None => raw_image,
-		SubtractBlack::Value(_) => todo!(),
-		SubtractBlack::CfaGrid(_) => subtract_black_cfa_grid(raw_image),
-	};
-
-	raw_image.black = SubtractBlack::None;
-	raw_image
-}
-
-pub fn subtract_black_cfa_grid(mut raw_image: RawImage) -> RawImage {
-	let width = raw_image.width;
-	let black_level = match raw_image.black {
-		SubtractBlack::CfaGrid(x) => x,
-		_ => unreachable!(),
-	};
-
-	for row in 0..raw_image.height {
-		for col in 0..width {
-			raw_image.data[row * width + col] = raw_image.data[row * width + col].saturating_sub(black_level[2 * (row % 2) + (col % 2)]);
+impl RawImage {
+	pub fn subtract_black_fn(&self) -> impl Fn(RawPixel) -> u16 {
+		match self.black {
+			SubtractBlack::CfaGrid(black_levels) => move |pixel: RawPixel| pixel.value.saturating_sub(black_levels[2 * (pixel.row % 2) + (pixel.column % 2)]),
+			_ => todo!(),
 		}
 	}
-
-	raw_image.maximum -= black_level.iter().max().unwrap();
-
-	raw_image
 }