Raw-rs: Add preprocessing and demosaicing steps (#1796)

* add subtract black step

* add scale colors step

* add raw to image step

* implement linear demosiacing and fix errors in previous code

* fix missing variable

* make dependencies of tests optional

* fix error in raw-rs tests

* fix typo in "demosiacing"

* use camera data from ADC and remove downloader

* cargo fmt

* use file_stem instead of file_name

* remove old camera data

* use equality instead of subtring to find model

* store camera_to_xyz in decimal form

* Code review

---------

Co-authored-by: Keavon Chambers <keavon@keavon.com>

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

@@ -1,7 +1,7 @@
 use crate::tiff::file::TiffRead;
 use crate::tiff::tags::SonyDataOffset;
 use crate::tiff::Ifd;
-use crate::RawImage;
+use crate::{RawImage, SubtractBlack};
 
 use bitstream_io::{BitRead, BitReader, Endianness, BE};
 use std::io::{Read, Seek};
@@ -22,6 +22,10 @@ pub fn decode_a100<R: Read + Seek>(ifd: Ifd, file: &mut TiffRead<R>) -> RawImage
 		data: image,
 		width: image_width,
 		height: image_height,
+		cfa_pattern: todo!(),
+		maximum: (1 << 12) - 1,
+		black: SubtractBlack::None,
+		camera_to_xyz: None,
 	}
 }
 

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

@@ -2,7 +2,7 @@ use crate::tiff::file::{Endian, TiffRead};
 use crate::tiff::tags::{BitsPerSample, CfaPattern, CfaPatternDim, Compression, ImageLength, ImageWidth, SonyToneCurve, StripByteCounts, StripOffsets, Tag};
 use crate::tiff::values::CurveLookupTable;
 use crate::tiff::{Ifd, TiffError};
-use crate::RawImage;
+use crate::{RawImage, SubtractBlack};
 
 use std::io::{Read, Seek};
 use tag_derive::Tag;
@@ -30,7 +30,9 @@ pub fn decode<R: Read + Seek>(ifd: Ifd, file: &mut TiffRead<R>) -> RawImage {
 
 	let image_width: usize = ifd.image_width.try_into().unwrap();
 	let image_height: usize = ifd.image_height.try_into().unwrap();
-	let _bits_per_sample: usize = ifd.bits_per_sample.into();
+	let bits_per_sample: usize = ifd.bits_per_sample.into();
+	assert!(bits_per_sample == 12);
+
 	let [cfa_pattern_width, cfa_pattern_height] = ifd.cfa_pattern_dim;
 	assert!(cfa_pattern_width == 2 && cfa_pattern_height == 2);
 
@@ -44,6 +46,10 @@ pub fn decode<R: Read + Seek>(ifd: Ifd, file: &mut TiffRead<R>) -> RawImage {
 		data: image,
 		width: image_width,
 		height: image_height,
+		cfa_pattern: ifd.cfa_pattern.try_into().unwrap(),
+		maximum: (1 << 14) - 1,
+		black: SubtractBlack::CfaGrid([512, 512, 512, 512]), // TODO: Find the correct way to do this
+		camera_to_xyz: None,
 	}
 }
 

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

@@ -1,7 +1,8 @@
 use crate::tiff::file::TiffRead;
-use crate::tiff::tags::{BitsPerSample, CfaPattern, CfaPatternDim, Compression, ImageLength, ImageWidth, RowsPerStrip, StripByteCounts, StripOffsets, Tag};
+use crate::tiff::tags::{BitsPerSample, BlackLevel, CfaPattern, CfaPatternDim, Compression, ImageLength, ImageWidth, RowsPerStrip, StripByteCounts, StripOffsets, Tag};
 use crate::tiff::{Ifd, TiffError};
-use crate::RawImage;
+use crate::{RawImage, SubtractBlack};
+
 use std::io::{Read, Seek};
 use tag_derive::Tag;
 
@@ -13,6 +14,7 @@ struct ArwUncompressedIfd {
 	rows_per_strip: RowsPerStrip,
 	bits_per_sample: BitsPerSample,
 	compression: Compression,
+	black_level: BlackLevel,
 	cfa_pattern: CfaPattern,
 	cfa_pattern_dim: CfaPatternDim,
 	strip_offsets: StripOffsets,
@@ -29,7 +31,7 @@ pub fn decode<R: Read + Seek>(ifd: Ifd, file: &mut TiffRead<R>) -> RawImage {
 	let image_width: usize = ifd.image_width.try_into().unwrap();
 	let image_height: usize = ifd.image_height.try_into().unwrap();
 	let rows_per_strip: usize = ifd.rows_per_strip.try_into().unwrap();
-	let _bits_per_sample: usize = ifd.bits_per_sample.into();
+	let bits_per_sample: usize = ifd.bits_per_sample.into();
 	let [cfa_pattern_width, cfa_pattern_height] = ifd.cfa_pattern_dim;
 	assert!(cfa_pattern_width == 2 && cfa_pattern_height == 2);
 
@@ -52,5 +54,9 @@ pub fn decode<R: Read + Seek>(ifd: Ifd, file: &mut TiffRead<R>) -> RawImage {
 		data: image,
 		width: image_width,
 		height: image_height,
+		cfa_pattern: ifd.cfa_pattern.try_into().unwrap(),
+		maximum: if bits_per_sample == 16 { u16::MAX } else { (1 << bits_per_sample) - 1 },
+		black: SubtractBlack::CfaGrid(ifd.black_level),
+		camera_to_xyz: None,
 	}
 }

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

@@ -0,0 +1,77 @@
+use crate::{Image, RawImage};
+
+fn average(data: &[u16], indexes: impl Iterator<Item = i64>) -> u16 {
+	let mut sum = 0;
+	let mut count = 0;
+	for index in indexes {
+		if index >= 0 && (index as usize) < data.len() {
+			sum += data[index as usize] as u32;
+			count += 1;
+		}
+	}
+
+	(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!(),
+	}
+}
+
+fn linear_demosaic_rggb(mut raw_image: RawImage) -> Image<u16> {
+	let width = raw_image.width as i64;
+	let height = raw_image.height as i64;
+
+	for row in 0..height {
+		let row_by_width = row * width;
+
+		for col in 0..width {
+			let pixel_index = row_by_width + col;
+
+			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];
+
+			match (row % 2 == 0, col % 2 == 0) {
+				(true, true) => {
+					let indexes = cross_indexes.iter().map(|x| 3 * x + 1);
+					raw_image.data[3 * (pixel_index as usize) + 1] = average(&raw_image.data, indexes);
+
+					let indexes = diagonal_indexes.iter().map(|x| 3 * x + 2);
+					raw_image.data[3 * (pixel_index as usize) + 2] = average(&raw_image.data, indexes);
+				}
+				(true, false) => {
+					let indexes = horizontal_indexes.iter().map(|x| 3 * x);
+					raw_image.data[3 * (pixel_index as usize)] = average(&raw_image.data, indexes);
+
+					let indexes = vertical_indexes.iter().map(|x| 3 * x + 2);
+					raw_image.data[3 * (pixel_index as usize) + 2] = average(&raw_image.data, indexes);
+				}
+				(false, true) => {
+					let indexes = vertical_indexes.iter().map(|x| 3 * x);
+					raw_image.data[3 * (pixel_index as usize)] = average(&raw_image.data, indexes);
+
+					let indexes = horizontal_indexes.iter().map(|x| 3 * x + 2);
+					raw_image.data[3 * (pixel_index as usize) + 2] = average(&raw_image.data, indexes);
+				}
+				(false, false) => {
+					let indexes = cross_indexes.iter().map(|x| 3 * x + 1);
+					raw_image.data[3 * (pixel_index as usize) + 1] = average(&raw_image.data, indexes);
+
+					let indexes = diagonal_indexes.iter().map(|x| 3 * x);
+					raw_image.data[3 * (pixel_index as usize)] = average(&raw_image.data, indexes);
+				}
+			}
+		}
+	}
+
+	Image {
+		channels: 3,
+		data: raw_image.data,
+		width: raw_image.width,
+		height: raw_image.height,
+	}
+}

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

@@ -0,0 +1 @@
+pub mod linear_demosaicing;

libraries/raw-rs/src/lib.rs

@@ -1,18 +1,33 @@
 pub mod decoder;
+pub mod demosaicing;
+pub mod metadata;
+pub mod preprocessing;
 pub mod tiff;
 
+use crate::preprocessing::camera_data::camera_to_xyz;
+
 use tag_derive::Tag;
 use tiff::file::TiffRead;
-use tiff::tags::{Compression, ImageLength, ImageWidth, Model, StripByteCounts, SubIfd, Tag};
+use tiff::tags::{Compression, ImageLength, ImageWidth, StripByteCounts, SubIfd, Tag};
 use tiff::{Ifd, TiffError};
 
 use std::io::{Read, Seek};
 use thiserror::Error;
 
+pub enum SubtractBlack {
+	None,
+	Value(u16),
+	CfaGrid([u16; 4]),
+}
+
 pub struct RawImage {
 	pub data: Vec<u16>,
 	pub width: usize,
 	pub height: usize,
+	pub cfa_pattern: [u8; 4],
+	pub maximum: u16,
+	pub black: SubtractBlack,
+	pub camera_to_xyz: Option<[f64; 9]>,
 }
 
 pub struct Image<T> {
@@ -35,28 +50,41 @@ pub fn decode<R: Read + Seek>(reader: &mut R) -> Result<RawImage, DecoderError>
 	let mut file = TiffRead::new(reader)?;
 	let ifd = Ifd::new_first_ifd(&mut file)?;
 
-	// TODO: This is only for the tests to pass for now. Replace this with the correct implementation when the decoder is complete.
-	let model = ifd.get_value::<Model, _>(&mut file)?;
+	let camera_model = metadata::identify::identify_camera_model(&ifd, &mut file).unwrap();
 
-	if model == "DSLR-A100" {
-		Ok(decoder::arw1::decode_a100(ifd, &mut file))
+	let mut raw_image = if camera_model.model == "DSLR-A100" {
+		decoder::arw1::decode_a100(ifd, &mut file)
 	} else {
 		let sub_ifd = ifd.get_value::<SubIfd, _>(&mut file)?;
 		let arw_ifd = sub_ifd.get_value::<ArwIfd, _>(&mut file)?;
 
 		if arw_ifd.compression == 1 {
-			Ok(decoder::uncompressed::decode(sub_ifd, &mut file))
+			decoder::uncompressed::decode(sub_ifd, &mut file)
 		} else if arw_ifd.strip_byte_counts[0] == arw_ifd.image_width * arw_ifd.image_height {
-			Ok(decoder::arw2::decode(sub_ifd, &mut file))
+			decoder::arw2::decode(sub_ifd, &mut file)
 		} else {
 			// TODO: implement for arw 1.
 			todo!()
 		}
-	}
+	};
+
+	raw_image.camera_to_xyz = camera_to_xyz(&camera_model);
+
+	Ok(raw_image)
 }
 
-pub fn process_8bit(_image: RawImage) -> Image<u8> {
-	todo!()
+pub fn process_8bit(raw_image: RawImage) -> Image<u8> {
+	let raw_image = crate::preprocessing::subtract_black::subtract_black(raw_image);
+	let raw_image = crate::preprocessing::raw_to_image::raw_to_image(raw_image);
+	let raw_image = crate::preprocessing::scale_colors::scale_colors(raw_image);
+	let image = crate::demosaicing::linear_demosaicing::linear_demosaic(raw_image);
+
+	Image {
+		channels: image.channels,
+		data: image.data.iter().map(|x| (x >> 8) as u8).collect(),
+		width: image.width,
+		height: image.height,
+	}
 }
 
 pub fn process_16bit(_image: RawImage) -> Image<u16> {

libraries/raw-rs/src/metadata/identify.rs

@@ -0,0 +1,60 @@
+use crate::tiff::file::TiffRead;
+use crate::tiff::tags::{Make, Model, Tag};
+use crate::tiff::{Ifd, TiffError};
+
+use std::io::{Read, Seek};
+use tag_derive::Tag;
+
+const COMPANY_NAMES: [&str; 22] = [
+	"AgfaPhoto",
+	"Canon",
+	"Casio",
+	"Epson",
+	"Fujifilm",
+	"Mamiya",
+	"Minolta",
+	"Motorola",
+	"Kodak",
+	"Konica",
+	"Leica",
+	"Nikon",
+	"Nokia",
+	"Olympus",
+	"Ricoh",
+	"Pentax",
+	"Phase One",
+	"Samsung",
+	"Sigma",
+	"Sinar",
+	"Sony",
+	"YI",
+];
+
+#[allow(dead_code)]
+#[derive(Tag)]
+struct CameraModelIfd {
+	make: Make,
+	model: Model,
+}
+
+pub struct CameraModel {
+	pub make: String,
+	pub model: String,
+}
+
+pub fn identify_camera_model<R: Read + Seek>(ifd: &Ifd, file: &mut TiffRead<R>) -> Option<CameraModel> {
+	let mut ifd = ifd.get_value::<CameraModelIfd, _>(file).unwrap();
+
+	ifd.make.make_ascii_lowercase();
+	for company_name in COMPANY_NAMES {
+		let lowercase_company_name = company_name.to_ascii_lowercase();
+		if ifd.make.contains(&lowercase_company_name) {
+			return Some(CameraModel {
+				make: company_name.to_string(),
+				model: ifd.model,
+			});
+		}
+	}
+
+	None
+}

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

@@ -0,0 +1 @@
+pub mod identify;

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

@@ -0,0 +1,26 @@
+use crate::metadata::identify::CameraModel;
+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!();
+
+pub fn camera_to_xyz(camera_model: &CameraModel) -> Option<[f64; 9]> {
+	let camera_name_needle = camera_model.make.to_owned() + " " + &camera_model.model;
+	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.))
+}

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

@@ -0,0 +1,4 @@
+pub mod camera_data;
+pub mod raw_to_image;
+pub mod scale_colors;
+pub mod subtract_black;

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

@@ -0,0 +1,17 @@
+use crate::RawImage;
+
+pub fn raw_to_image(mut raw_image: RawImage) -> RawImage {
+	let mut image = Vec::with_capacity(raw_image.width * raw_image.height * 3);
+
+	for row in 0..raw_image.height {
+		for col in 0..raw_image.width {
+			let mut pixel = [0_u16; 3];
+			let color_index = raw_image.cfa_pattern[2 * (row % 2) + (col % 2)];
+			pixel[color_index as usize] = raw_image.data[row * raw_image.width + col];
+			image.extend_from_slice(&pixel);
+		}
+	}
+
+	raw_image.data = image;
+	raw_image
+}

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

@@ -0,0 +1,46 @@
+use crate::RawImage;
+
+const XYZ_TO_RGB: [[f64; 3]; 3] = [[0.412453, 0.357580, 0.180423], [0.212671, 0.715160, 0.072169], [0.019334, 0.119193, 0.950227]];
+
+pub fn scale_colors(mut raw_image: RawImage) -> RawImage {
+	if let Some(camera_to_xyz) = raw_image.camera_to_xyz {
+		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 mut white_balance_multiplier = camera_to_rgb.map(|x| 1. / x.iter().sum::<f64>());
+
+		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().fold(f64::INFINITY, |a, &b| a.min(b))
+		} else {
+			white_balance_multiplier.iter().fold(f64::NEG_INFINITY, |a, &b| a.max(b))
+		};
+
+		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.]
+		};
+
+		for i in 0..(raw_image.height * raw_image.width) {
+			for (c, multiplier) in final_multiplier.iter().enumerate() {
+				raw_image.data[3 * i + c] = ((raw_image.data[3 * i + c] as f64) * multiplier).min(u16::MAX as f64).max(0.) as u16;
+			}
+		}
+	}
+
+	raw_image
+}

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

@@ -0,0 +1,30 @@
+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)]);
+		}
+	}
+
+	raw_image.maximum -= black_level.iter().max().unwrap();
+
+	raw_image
+}

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

@@ -29,8 +29,11 @@ pub enum TagId {
 	JpegOffset = 0x201,
 	JpegLength = 0x202,
 	SonyToneCurve = 0x7010,
+	BlackLevel = 0x7310,
 	CfaPatternDim = 0x828d,
 	CfaPattern = 0x828e,
+	ColorMatrix1 = 0xc621,
+	ColorMatrix2 = 0xc622,
 
 	#[num_enum(catch_all)]
 	Unknown(u16),

libraries/raw-rs/src/tiff/tags.rs

@@ -1,4 +1,4 @@
-use super::types::{Array, ConstArray, TagType, TypeByte, TypeIfd, TypeLong, TypeNumber, TypeShort, TypeSonyToneCurve, TypeString};
+use super::types::{Array, ConstArray, TagType, TypeByte, TypeIfd, TypeLong, TypeNumber, TypeSRational, TypeShort, TypeSonyToneCurve, TypeString};
 use super::{Ifd, TagId, TiffError, TiffRead};
 
 use std::io::{Read, Seek};
@@ -24,10 +24,13 @@ pub struct StripByteCounts;
 pub struct SubIfd;
 pub struct JpegOffset;
 pub struct JpegLength;
-pub struct CfaPatternDim;
-pub struct CfaPattern;
 pub struct SonyDataOffset;
 pub struct SonyToneCurve;
+pub struct BlackLevel;
+pub struct CfaPatternDim;
+pub struct CfaPattern;
+pub struct ColorMatrix1;
+pub struct ColorMatrix2;
 
 impl SimpleTag for ImageWidth {
 	type Type = TypeNumber;
@@ -141,6 +144,20 @@ impl SimpleTag for CfaPattern {
 	const NAME: &'static str = "CFA Pattern";
 }
 
+impl SimpleTag for ColorMatrix1 {
+	type Type = Array<TypeSRational>;
+
+	const ID: TagId = TagId::ColorMatrix1;
+	const NAME: &'static str = "Color Matrix 1";
+}
+
+impl SimpleTag for ColorMatrix2 {
+	type Type = Array<TypeSRational>;
+
+	const ID: TagId = TagId::ColorMatrix2;
+	const NAME: &'static str = "Color Matrix 2";
+}
+
 impl SimpleTag for SonyDataOffset {
 	type Type = TypeLong;
 
@@ -155,6 +172,12 @@ impl SimpleTag for SonyToneCurve {
 	const NAME: &'static str = "Sony Tone Curve";
 }
 
+impl SimpleTag for BlackLevel {
+	const ID: TagId = TagId::BlackLevel;
+	type Type = ConstArray<TypeShort, 4>;
+	const NAME: &'static str = "Black Level";
+}
+
 pub trait Tag {
 	type Output;
 

libraries/raw-rs/src/tiff/values.rs

@@ -1,8 +1,30 @@
-pub struct Rational<T> {
+pub trait ToFloat {
+	fn to_float(&self) -> f64;
+}
+
+impl ToFloat for u32 {
+	fn to_float(&self) -> f64 {
+		*self as f64
+	}
+}
+
+impl ToFloat for i32 {
+	fn to_float(&self) -> f64 {
+		*self as f64
+	}
+}
+
+pub struct Rational<T: ToFloat> {
 	pub numerator: T,
 	pub denominator: T,
 }
 
+impl<T: ToFloat> ToFloat for Rational<T> {
+	fn to_float(&self) -> f64 {
+		self.numerator.to_float() / self.denominator.to_float()
+	}
+}
+
 pub struct CurveLookupTable {
 	table: Vec<u16>,
 }