From f685b5e87878306ac3175fefaf6f31c7809612fb Mon Sep 17 00:00:00 2001 From: Elbert Ronnie Date: Sun, 13 Oct 2024 00:09:46 +0530 Subject: [PATCH] Include transform and gamma correction step --- .../src/demosaicing/linear_demosaicing.rs | 40 +++++----- libraries/raw-rs/src/lib.rs | 78 ++++++++++++++----- .../src/postprocessing/convert_to_rgb.rs | 27 +++---- .../src/postprocessing/gamma_correction.rs | 37 ++++----- .../raw-rs/src/postprocessing/transform.rs | 68 ++++++++-------- libraries/raw-rs/src/processing.rs | 20 ++--- 6 files changed, 152 insertions(+), 118 deletions(-) diff --git a/libraries/raw-rs/src/demosaicing/linear_demosaicing.rs b/libraries/raw-rs/src/demosaicing/linear_demosaicing.rs index e4f998e9bf..1a5b4edb83 100644 --- a/libraries/raw-rs/src/demosaicing/linear_demosaicing.rs +++ b/libraries/raw-rs/src/demosaicing/linear_demosaicing.rs @@ -1,4 +1,4 @@ -use crate::{Pixel, RawImage}; +use crate::{Pixel, RawImage ,Captures}; fn average(data: &[u16], indexes: impl Iterator) -> u16 { let mut sum = 0; @@ -13,12 +13,6 @@ fn average(data: &[u16], indexes: impl Iterator) -> u16 { (sum / count) as u16 } -// This trait is here only to circumvent Rust's lifetime capturing rules in return type impl Trait. -// See https://youtu.be/CWiz_RtA1Hw?si=j0si4qE2Y20f71Uo -// This should be removed when Rust 2024 edition is released as described in https://blog.rust-lang.org/2024/09/05/impl-trait-capture-rules.html -pub trait Captures {} -impl Captures for T {} - impl RawImage { pub fn linear_demosaic_iter(&self) -> impl Iterator + Captures<&'_ ()> { match self.cfa_pattern { @@ -45,30 +39,38 @@ impl RawImage { let pixel_index = pixel_index as usize; match (row % 2 == 0, column % 2 == 0) { (true, true) => Pixel { - red: self.data[pixel_index], - blue: average(&self.data, cross_indexes.into_iter()), - green: average(&self.data, diagonal_indexes.into_iter()), + 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 { - red: average(&self.data, horizontal_indexes.into_iter()), - blue: self.data[pixel_index], - green: average(&self.data, vertical_indexes.into_iter()), + 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 { - red: average(&self.data, vertical_indexes.into_iter()), - blue: self.data[pixel_index], - green: average(&self.data, horizontal_indexes.into_iter()), + 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 { - red: average(&self.data, diagonal_indexes.into_iter()), - blue: average(&self.data, cross_indexes.into_iter()), - green: self.data[pixel_index], + 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, }, diff --git a/libraries/raw-rs/src/lib.rs b/libraries/raw-rs/src/lib.rs index e9574f11b0..b52e94d998 100644 --- a/libraries/raw-rs/src/lib.rs +++ b/libraries/raw-rs/src/lib.rs @@ -18,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), @@ -47,8 +50,6 @@ pub struct Image { /// See 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)] @@ -98,29 +99,30 @@ pub fn process_8bit(raw_image: RawImage) -> Image { 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 { - let mut raw_image = crate::preprocessing::camera_data::calculate_conversion_matrices(raw_image); + let raw_image = crate::preprocessing::camera_data::calculate_conversion_matrices(raw_image); + let subtract_black = raw_image.subtract_black_fn(); let scale_colors = raw_image.scale_colors_fn(); + let raw_image = raw_image.apply((subtract_black, scale_colors)); + 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)); - raw_image.apply((subtract_black, scale_colors)); - let mut image = raw_image.demosaic_and_apply((convert_to_rgb, &mut record_histogram)); - - image.histogram = Some(record_histogram.histogram); - - let image = crate::postprocessing::transform::transform(image); - crate::postprocessing::gamma_correction::gamma_correction(image) + 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 { - fn apply(&mut self, mut transform: impl RawPixelTransform) { + pub fn apply(mut self, mut transform: impl RawPixelTransform) -> RawImage { for (index, value) in self.data.iter_mut().enumerate() { let pixel = RawPixel { value: *value, @@ -129,15 +131,15 @@ impl RawImage { }; *value = transform.apply(pixel); } + + self } - fn demosaic_and_apply(self, mut transform: impl PixelTransform) -> Image { + pub fn demosaic_and_apply(self, mut transform: impl PixelTransform) -> Image { let mut image = vec![0; self.width * self.height * 3]; - for Pixel { red, blue, green, row, column } in self.linear_demosaic_iter().map(|pixel| transform.apply(pixel)) { + 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] = red; - image[3 * pixel_index + 1] = blue; - image[3 * pixel_index + 2] = green; + image[3 * pixel_index..3 * (pixel_index+1)].copy_from_slice(&values); } Image { @@ -146,8 +148,38 @@ impl RawImage { width: self.width, height: self.height, transform: self.transform, - rgb_to_camera: self.rgb_to_camera, - histogram: None, + } + } +} + +impl Image { + pub fn apply(mut self, mut transform: impl PixelTransform) -> Image { + 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 { + 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, } } } @@ -161,3 +193,9 @@ pub enum DecoderError { #[error("An IO Error ocurred")] IoError(#[from] std::io::Error), } + +// This trait is here only to circumvent Rust's lifetime capturing rules in return type impl Trait. +// See https://youtu.be/CWiz_RtA1Hw?si=j0si4qE2Y20f71Uo +// This should be removed when Rust 2024 edition is released as described in https://blog.rust-lang.org/2024/09/05/impl-trait-capture-rules.html +pub trait Captures {} +impl Captures for T {} diff --git a/libraries/raw-rs/src/postprocessing/convert_to_rgb.rs b/libraries/raw-rs/src/postprocessing/convert_to_rgb.rs index 278e93942e..23fd7d8f66 100644 --- a/libraries/raw-rs/src/postprocessing/convert_to_rgb.rs +++ b/libraries/raw-rs/src/postprocessing/convert_to_rgb.rs @@ -1,18 +1,13 @@ -use crate::{Pixel, PixelTransform, RawImage}; - -const CHANNELS_IN_RGB: usize = 3; +use crate::{Pixel, PixelTransform, RawImage, CHANNELS_IN_RGB, Histogram}; impl RawImage { - pub fn convert_to_rgb_fn(&self) -> impl Fn(Pixel) -> Pixel { + pub fn convert_to_rgb_fn(&self) -> impl Fn(Pixel) -> [u16; CHANNELS_IN_RGB] { let Some(rgb_to_camera) = self.rgb_to_camera else { todo!() }; - move |mut pixel: Pixel| { - let [red, blue, green] = std::array::from_fn(|i| i).map(|i| rgb_to_camera[i][0] * pixel.red as f64 + rgb_to_camera[i][1] * pixel.blue as f64 + rgb_to_camera[i][2] * pixel.green as f64); - - pixel.red = (red as u16).clamp(0, u16::MAX); - pixel.green = (green as u16).clamp(0, u16::MAX); - pixel.blue = (blue as u16).clamp(0, u16::MAX); - pixel + move |pixel: Pixel| { + std::array::from_fn(|i| i) + .map(|i| rgb_to_camera[i].iter().zip(pixel.values.iter()).map(|(&coeff, &value)| coeff * value as f64).sum()) + .map(|x: f64| (x as u16).clamp(0, u16::MAX)) } } @@ -22,7 +17,7 @@ impl RawImage { } pub struct RecordHistogram { - pub histogram: [[usize; 0x2000]; CHANNELS_IN_RGB], + pub histogram: Histogram, } impl RecordHistogram { @@ -34,10 +29,8 @@ impl RecordHistogram { } impl PixelTransform for &mut RecordHistogram { - fn apply(&mut self, pixel: Pixel) -> Pixel { - self.histogram[0][pixel.red as usize >> CHANNELS_IN_RGB] += 1; - self.histogram[1][pixel.blue as usize >> CHANNELS_IN_RGB] += 1; - self.histogram[2][pixel.green as usize >> CHANNELS_IN_RGB] += 1; - pixel + fn apply(&mut self, pixel: Pixel) -> [u16; CHANNELS_IN_RGB] { + self.histogram.iter_mut().zip(pixel.values.iter()).for_each(|(histogram, &value)| histogram[value as usize >> CHANNELS_IN_RGB] += 1); + pixel.values } } diff --git a/libraries/raw-rs/src/postprocessing/gamma_correction.rs b/libraries/raw-rs/src/postprocessing/gamma_correction.rs index e3b98b2174..690c73b84a 100644 --- a/libraries/raw-rs/src/postprocessing/gamma_correction.rs +++ b/libraries/raw-rs/src/postprocessing/gamma_correction.rs @@ -1,32 +1,29 @@ -use crate::Image; +use crate::{Image, Pixel, Histogram, CHANNELS_IN_RGB}; use std::f64::consts::E; -pub fn gamma_correction(mut image: Image) -> Image { - let Some(histogram) = image.histogram else { return image }; +impl Image { + pub fn gamma_correction_fn(&self, histogram: &Histogram) -> impl Fn(Pixel) -> [u16; CHANNELS_IN_RGB] { + let percentage = self.width * self.height; - let percentage = image.width * image.height; + let mut white = 0; + for channel_histogram in histogram { + let mut total = 0; + for i in (0x20..0x2000).rev() { + total += channel_histogram[i] as u64; - let mut white = 0; - for channel_histogram in histogram { - let mut total = 0; - for i in (0x20..0x2000).rev() { - total += channel_histogram[i] as u64; - - if total * 100 > percentage as u64 { - white = white.max(i); - break; + if total * 100 > percentage as u64 { + white = white.max(i); + break; + } } } - } - let curve = generate_gamma_curve(0.45, 4.5, (white << 3) as f64); + let curve = generate_gamma_curve(0.45, 4.5, (white << 3) as f64); - for value in image.data.iter_mut() { - *value = curve[*value as usize]; + move |pixel: Pixel| { + pixel.values.map(|value| curve[value as usize]) + } } - image.histogram = None; - - image } /// `max_intensity` must be non-zero. diff --git a/libraries/raw-rs/src/postprocessing/transform.rs b/libraries/raw-rs/src/postprocessing/transform.rs index 49f9a58104..d91d3bf309 100644 --- a/libraries/raw-rs/src/postprocessing/transform.rs +++ b/libraries/raw-rs/src/postprocessing/transform.rs @@ -1,44 +1,48 @@ -use crate::{Image, Transform}; +use crate::{Image, Transform, Captures, Pixel}; -pub fn transform(mut image: Image) -> Image { - if image.transform.is_identity() { - return image; - } +impl Image { + pub fn transform_iter(&self) -> (usize, usize, impl Iterator + Captures<&'_ ()>) { + let (final_width, final_height) = if self.transform.will_swap_coordinates() { + (self.height, self.width) + } else { + (self.width, self.height) + }; - let channels = image.channels as usize; - let mut data = vec![0; channels * image.width * image.height]; + let index_0_0 = inverse_transform_index(self.transform, 0, 0, self.width, self.height); + let index_0_1 = inverse_transform_index(self.transform, 0, 1, self.width, self.height); + let index_1_0 = inverse_transform_index(self.transform, 1, 0, self.width, self.height); - let (final_width, final_height) = if image.transform.will_swap_coordinates() { - (image.height, image.width) - } else { - (image.width, image.height) - }; + let column_step = (index_0_1.0 - index_0_0.0, index_0_1.1 - index_0_0.1); + let row_step = (index_1_0.0 - index_0_0.0, index_1_0.1 - index_0_0.1); + let mut index = index_0_0; - let mut initial_index = inverse_transform_index(image.transform, 0, 0, image.width, image.height); - let column_step = inverse_transform_index(image.transform, 0, 1, image.width, image.height) as i64 - initial_index as i64; - let row_step = inverse_transform_index(image.transform, 1, 0, image.width, image.height) as i64 - inverse_transform_index(image.transform, 0, final_width, image.width, image.height) as i64; + let channels = self.channels as usize; - for row in 0..final_height { - for col in 0..final_width { - let transformed_index = final_width * row + col; + ( + final_width, + final_height, + (0..final_height).flat_map(move |row| { + let temp = (0..final_width).map(move |column| { + let initial_index = (self.width as i64 * index.0 + index.1) as usize; + let pixel = &self.data[channels * initial_index..channels * (initial_index + 1)]; + index = (index.0 + column_step.0, index.1 + column_step.1); - let copy_from_range = channels * initial_index..channels * (initial_index + 1); - let copy_to_range = channels * transformed_index..channels * (transformed_index + 1); - data[copy_to_range].copy_from_slice(&image.data[copy_from_range]); + Pixel { + values: pixel.try_into().unwrap(), + row, + column, + } + }); - initial_index = (initial_index as i64 + column_step) as usize; - } - initial_index = (initial_index as i64 + row_step) as usize; - } + index = (index.0 + row_step.0, index.1 + row_step.1); - image.data = data; - image.width = final_width; - image.height = final_height; - - image + temp + }) + ) + } } -pub fn inverse_transform_index(transform: Transform, mut row: usize, mut column: usize, width: usize, height: usize) -> usize { +pub fn inverse_transform_index(transform: Transform, mut row: usize, mut column: usize, width: usize, height: usize) -> (i64, i64) { let value = match transform { Transform::Horizontal => 0, Transform::MirrorHorizontal => 1, @@ -62,5 +66,5 @@ pub fn inverse_transform_index(transform: Transform, mut row: usize, mut column: column = width - 1 - column; } - width * row + column + (row as i64, column as i64) } diff --git a/libraries/raw-rs/src/processing.rs b/libraries/raw-rs/src/processing.rs index 7f94685c1b..d68325595e 100644 --- a/libraries/raw-rs/src/processing.rs +++ b/libraries/raw-rs/src/processing.rs @@ -1,3 +1,5 @@ +use crate::CHANNELS_IN_RGB; + #[derive(Clone, Copy)] pub struct RawPixel { pub value: u16, @@ -7,9 +9,7 @@ pub struct RawPixel { #[derive(Clone, Copy)] pub struct Pixel { - pub red: u16, - pub blue: u16, - pub green: u16, + pub values: [u16; CHANNELS_IN_RGB], pub row: usize, pub column: usize, } @@ -34,20 +34,20 @@ impl RawPixelTransform for (T1, T2 } pub trait PixelTransform { - fn apply(&mut self, pixel: Pixel) -> Pixel; + fn apply(&mut self, pixel: Pixel) -> [u16; CHANNELS_IN_RGB]; } -impl Pixel> PixelTransform for T { - fn apply(&mut self, pixel: Pixel) -> Pixel { +impl [u16; CHANNELS_IN_RGB]> PixelTransform for T { + fn apply(&mut self, pixel: Pixel) -> [u16; CHANNELS_IN_RGB] { self(pixel) } } impl PixelTransform for (T1, T2) { - fn apply(&mut self, mut pixel: Pixel) -> Pixel { - pixel = self.0.apply(pixel); - pixel = self.1.apply(pixel); + fn apply(&mut self, mut pixel: Pixel) -> [u16; CHANNELS_IN_RGB] { + pixel.values = self.0.apply(pixel); + pixel.values = self.1.apply(pixel); - pixel + pixel.values } }