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/Cargo.toml

@@ -26,6 +26,7 @@ thiserror = { workspace = true }
 # Required dependencies
 bitstream-io = "2.3.0"
 num_enum = "0.7.2"
+fortuples = "0.9.1"
 
 # Optional workspace dependencies
 image = { workspace = true, optional = true }

libraries/raw-rs/src/decoder/arw1.rs

@@ -28,8 +28,8 @@ pub fn decode_a100<R: Read + Seek>(ifd: Ifd, file: &mut TiffRead<R>) -> RawImage
 		black: SubtractBlack::None,
 		transform: Transform::Horizontal,
 		camera_model: None,
-		camera_white_balance_multiplier: None,
-		white_balance_multiplier: None,
+		camera_white_balance: None,
+		white_balance: None,
 		camera_to_rgb: None,
 		rgb_to_camera: None,
 	}

libraries/raw-rs/src/decoder/arw2.rs

@@ -52,8 +52,8 @@ pub fn decode<R: Read + Seek>(ifd: Ifd, file: &mut TiffRead<R>) -> RawImage {
 		black: SubtractBlack::CfaGrid([512, 512, 512, 512]), // TODO: Find the correct way to do this
 		transform: Transform::Horizontal,
 		camera_model: None,
-		camera_white_balance_multiplier: ifd.white_balance_levels.map(|arr| arr.map(|x| x as f64)),
-		white_balance_multiplier: None,
+		camera_white_balance: ifd.white_balance_levels.map(|arr| arr.map(|x| x as f64)),
+		white_balance: None,
 		camera_to_rgb: None,
 		rgb_to_camera: None,
 	}

libraries/raw-rs/src/decoder/uncompressed.rs

@@ -60,8 +60,8 @@ pub fn decode<R: Read + Seek>(ifd: Ifd, file: &mut TiffRead<R>) -> RawImage {
 		black: SubtractBlack::CfaGrid(ifd.black_level),
 		transform: Transform::Horizontal,
 		camera_model: None,
-		camera_white_balance_multiplier: ifd.white_balance_levels.map(|arr| arr.map(|x| x as f64)),
-		white_balance_multiplier: None,
+		camera_white_balance: ifd.white_balance_levels.map(|arr| arr.map(|x| x as f64)),
+		white_balance: None,
 		camera_to_rgb: None,
 		rgb_to_camera: None,
 	}

libraries/raw-rs/src/demosaicing/linear_demosaicing.rs

@@ -1,4 +1,4 @@
-use crate::{Image, RawImage};
+use crate::{Pixel, RawImage};
 
 fn average(data: &[u16], indexes: impl Iterator<Item = i64>) -> u16 {
 	let mut sum = 0;
@@ -13,62 +13,69 @@ fn average(data: &[u16], indexes: impl Iterator<Item = i64>) -> u16 {
 	(sum / count) as u16
 }
 
-pub fn linear_demosaic(raw_image: RawImage) -> Image<u16> {
-	match raw_image.cfa_pattern {
-		[0, 1, 1, 2] => linear_demosaic_rggb(raw_image),
-		_ => todo!(),
+impl RawImage {
+	pub fn linear_demosaic_iter(&self) -> impl Iterator<Item = Pixel> + use<'_> {
+		match self.cfa_pattern {
+			[0, 1, 1, 2] => self.linear_demosaic_rggb_iter(),
+			_ => todo!(),
+		}
 	}
-}
 
-fn linear_demosaic_rggb(raw_image: RawImage) -> Image<u16> {
-	let mut image = vec![0; raw_image.width * raw_image.height * 3];
-	let width = raw_image.width as i64;
-	let height = raw_image.height as i64;
+	fn linear_demosaic_rggb_iter(&self) -> impl Iterator<Item = Pixel> + use<'_> {
+		let width = self.width as i64;
+		let height = self.height as i64;
 
-	for row in 0..height {
-		let row_by_width = row * width;
+		(0..height).flat_map(move |row| {
+			let row_by_width = row * width;
 
-		for col in 0..width {
-			let pixel_index = row_by_width + col;
+			(0..width).map(move |column| {
+				let pixel_index = row_by_width + column;
 
-			let vertical_indexes = [pixel_index + width, pixel_index - width];
-			let horizontal_indexes = [pixel_index + 1, pixel_index - 1];
-			let cross_indexes = [pixel_index + width, pixel_index - width, pixel_index + 1, pixel_index - 1];
-			let diagonal_indexes = [pixel_index + width + 1, pixel_index - width + 1, pixel_index + width - 1, pixel_index - width - 1];
+				let vertical_indexes = [pixel_index + width, pixel_index - width];
+				let horizontal_indexes = [pixel_index + 1, pixel_index - 1];
+				let cross_indexes = [pixel_index + width, pixel_index - width, pixel_index + 1, pixel_index - 1];
+				let diagonal_indexes = [pixel_index + width + 1, pixel_index - width + 1, pixel_index + width - 1, pixel_index - width - 1];
 
-			let pixel_index = pixel_index as usize;
-			match (row % 2 == 0, col % 2 == 0) {
-				(true, true) => {
-					image[3 * pixel_index] = raw_image.data[pixel_index];
-					image[3 * pixel_index + 1] = average(&raw_image.data, cross_indexes.into_iter());
-					image[3 * pixel_index + 2] = average(&raw_image.data, diagonal_indexes.into_iter());
-				}
-				(true, false) => {
-					image[3 * pixel_index] = average(&raw_image.data, horizontal_indexes.into_iter());
-					image[3 * pixel_index + 1] = raw_image.data[pixel_index];
-					image[3 * pixel_index + 2] = average(&raw_image.data, vertical_indexes.into_iter());
-				}
-				(false, true) => {
-					image[3 * pixel_index] = average(&raw_image.data, vertical_indexes.into_iter());
-					image[3 * pixel_index + 1] = raw_image.data[pixel_index];
-					image[3 * pixel_index + 2] = average(&raw_image.data, horizontal_indexes.into_iter());
+				let pixel_index = pixel_index as usize;
+				match (row % 2 == 0, column % 2 == 0) {
+					(true, true) => Pixel {
+						values: [
+							self.data[pixel_index],
+							average(&self.data, cross_indexes.into_iter()),
+							average(&self.data, diagonal_indexes.into_iter()),
+						],
+						row: row as usize,
+						column: column as usize,
+					},
+					(true, false) => Pixel {
+						values: [
+							average(&self.data, horizontal_indexes.into_iter()),
+							self.data[pixel_index],
+							average(&self.data, vertical_indexes.into_iter()),
+						],
+						row: row as usize,
+						column: column as usize,
+					},
+					(false, true) => Pixel {
+						values: [
+							average(&self.data, vertical_indexes.into_iter()),
+							self.data[pixel_index],
+							average(&self.data, horizontal_indexes.into_iter()),
+						],
+						row: row as usize,
+						column: column as usize,
+					},
+					(false, false) => Pixel {
+						values: [
+							average(&self.data, diagonal_indexes.into_iter()),
+							average(&self.data, cross_indexes.into_iter()),
+							self.data[pixel_index],
+						],
+						row: row as usize,
+						column: column as usize,
+					},
 				}
-				(false, false) => {
-					image[3 * pixel_index] = average(&raw_image.data, diagonal_indexes.into_iter());
-					image[3 * pixel_index + 1] = average(&raw_image.data, cross_indexes.into_iter());
-					image[3 * pixel_index + 2] = raw_image.data[pixel_index];
-				}
-			}
-		}
-	}
-
-	Image {
-		channels: 3,
-		data: image,
-		width: raw_image.width,
-		height: raw_image.height,
-		transform: raw_image.transform,
-		rgb_to_camera: raw_image.rgb_to_camera,
-		histogram: None,
+			})
+		})
 	}
 }

libraries/raw-rs/src/lib.rs

@@ -3,10 +3,12 @@ pub mod demosaicing;
 pub mod metadata;
 pub mod postprocessing;
 pub mod preprocessing;
+pub mod processing;
 pub mod tiff;
 
 use crate::metadata::identify::CameraModel;
 
+use processing::{Pixel, PixelTransform, RawPixel, RawPixelTransform};
 use tag_derive::Tag;
 use tiff::file::TiffRead;
 use tiff::tags::{Compression, ImageLength, ImageWidth, Orientation, StripByteCounts, SubIfd, Tag};
@@ -16,6 +18,9 @@ use tiff::{Ifd, TiffError};
 use std::io::{Read, Seek};
 use thiserror::Error;
 
+pub const CHANNELS_IN_RGB: usize = 3;
+pub type Histogram = [[usize; 0x2000]; CHANNELS_IN_RGB];
+
 pub enum SubtractBlack {
 	None,
 	Value(u16),
@@ -31,8 +36,8 @@ pub struct RawImage {
 	pub maximum: u16,
 	pub black: SubtractBlack,
 	pub camera_model: Option<CameraModel>,
-	pub camera_white_balance_multiplier: Option<[f64; 4]>,
-	pub white_balance_multiplier: Option<[f64; 4]>,
+	pub camera_white_balance: Option<[f64; 4]>,
+	pub white_balance: Option<[f64; 4]>,
 	pub camera_to_rgb: Option<[[f64; 3]; 3]>,
 	pub rgb_to_camera: Option<[[f64; 3]; 3]>,
 }
@@ -45,8 +50,6 @@ pub struct Image<T> {
 	/// See <https://github.com/GraphiteEditor/Graphite/pull/1923#discussion_r1725070342> for more information.
 	pub channels: u8,
 	pub transform: Transform,
-	pub rgb_to_camera: Option<[[f64; 3]; 3]>,
-	pub(crate) histogram: Option<[[usize; 0x2000]; 3]>,
 }
 
 #[allow(dead_code)]
@@ -84,6 +87,8 @@ pub fn decode<R: Read + Seek>(reader: &mut R) -> Result<RawImage, DecoderError>
 	raw_image.camera_model = Some(camera_model);
 	raw_image.transform = transform;
 
+	raw_image.calculate_conversion_matrices();
+
 	Ok(raw_image)
 }
 
@@ -96,19 +101,94 @@ pub fn process_8bit(raw_image: RawImage) -> Image<u8> {
 		width: image.width,
 		height: image.height,
 		transform: image.transform,
-		rgb_to_camera: image.rgb_to_camera,
-		histogram: image.histogram,
 	}
 }
 
 pub fn process_16bit(raw_image: RawImage) -> Image<u16> {
-	let raw_image = crate::preprocessing::camera_data::calculate_conversion_matrices(raw_image);
-	let raw_image = crate::preprocessing::subtract_black::subtract_black(raw_image);
-	let raw_image = crate::preprocessing::scale_colors::scale_colors(raw_image);
-	let image = crate::demosaicing::linear_demosaicing::linear_demosaic(raw_image);
-	let image = crate::postprocessing::convert_to_rgb::convert_to_rgb(image);
-	let image = crate::postprocessing::transform::transform(image);
-	crate::postprocessing::gamma_correction::gamma_correction(image)
+	let subtract_black = raw_image.subtract_black_fn();
+	let scale_white_balance = raw_image.scale_white_balance_fn();
+	let scale_to_16bit = raw_image.scale_to_16bit_fn();
+	let raw_image = raw_image.apply((subtract_black, scale_white_balance, scale_to_16bit));
+
+	let convert_to_rgb = raw_image.convert_to_rgb_fn();
+	let mut record_histogram = raw_image.record_histogram_fn();
+	let image = raw_image.demosaic_and_apply((convert_to_rgb, &mut record_histogram));
+
+	let gamma_correction = image.gamma_correction_fn(&record_histogram.histogram);
+	if image.transform == Transform::Horizontal {
+		image.apply(gamma_correction)
+	} else {
+		image.transform_and_apply(gamma_correction)
+	}
+}
+
+impl RawImage {
+	pub fn apply(mut self, mut transform: impl RawPixelTransform) -> RawImage {
+		for (index, value) in self.data.iter_mut().enumerate() {
+			let pixel = RawPixel {
+				value: *value,
+				row: index / self.width,
+				column: index % self.width,
+			};
+			*value = transform.apply(pixel);
+		}
+
+		self
+	}
+
+	pub fn demosaic_and_apply(self, mut transform: impl PixelTransform) -> Image<u16> {
+		let mut image = vec![0; self.width * self.height * 3];
+		for Pixel { values, row, column } in self.linear_demosaic_iter().map(|mut pixel| {
+			pixel.values = transform.apply(pixel);
+			pixel
+		}) {
+			let pixel_index = row * self.width + column;
+			image[3 * pixel_index..3 * (pixel_index + 1)].copy_from_slice(&values);
+		}
+
+		Image {
+			channels: 3,
+			data: image,
+			width: self.width,
+			height: self.height,
+			transform: self.transform,
+		}
+	}
+}
+
+impl Image<u16> {
+	pub fn apply(mut self, mut transform: impl PixelTransform) -> Image<u16> {
+		for (index, values) in self.data.chunks_exact_mut(3).enumerate() {
+			let pixel = Pixel {
+				values: values.try_into().unwrap(),
+				row: index / self.width,
+				column: index % self.width,
+			};
+			values.copy_from_slice(&transform.apply(pixel));
+		}
+
+		self
+	}
+
+	pub fn transform_and_apply(self, mut transform: impl PixelTransform) -> Image<u16> {
+		let mut image = vec![0; self.width * self.height * 3];
+		let (width, height, iter) = self.transform_iter();
+		for Pixel { values, row, column } in iter.map(|mut pixel| {
+			pixel.values = transform.apply(pixel);
+			pixel
+		}) {
+			let pixel_index = row * width + column;
+			image[3 * pixel_index..3 * (pixel_index + 1)].copy_from_slice(&values);
+		}
+
+		Image {
+			channels: 3,
+			data: image,
+			width,
+			height,
+			transform: Transform::Horizontal,
+		}
+	}
 }
 
 #[derive(Error, Debug)]