use alloc::vec;
use alloc::vec::Vec;
use archmage::prelude::*;
const R: usize = 0;
const G: usize = 1;
const B: usize = 2;
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
#[non_exhaustive]
pub enum DemosaicMethod {
Bilinear,
#[default]
MalvarHeCutler,
}
#[cfg(feature = "rawloader")]
pub fn demosaic_to_rgb_f32(
data: &[f32],
width: usize,
height: usize,
cfa: &rawloader::CFA,
method: DemosaicMethod,
) -> Vec<f32> {
match method {
DemosaicMethod::Bilinear => demosaic_bilinear(data, width, height, cfa),
DemosaicMethod::MalvarHeCutler => demosaic_malvar(data, width, height, cfa),
}
}
#[cfg(feature = "rawloader")]
fn demosaic_bilinear(data: &[f32], width: usize, height: usize, cfa: &rawloader::CFA) -> Vec<f32> {
let mut rgb = vec![0.0f32; width * height * 3];
for row in 0..height {
for col in 0..width {
let color = cfa.color_at(row, col);
let val = data[row * width + col];
let out_idx = (row * width + col) * 3;
match color {
R => {
rgb[out_idx] = val;
rgb[out_idx + 1] = green_at_rb_bilinear(data, width, height, row, col);
rgb[out_idx + 2] = opposite_at_rb_bilinear(data, width, height, row, col, cfa);
}
G => {
let (r, b) = rb_at_green_bilinear(data, width, height, row, col, cfa);
rgb[out_idx] = r;
rgb[out_idx + 1] = val;
rgb[out_idx + 2] = b;
}
B => {
rgb[out_idx] = opposite_at_rb_bilinear(data, width, height, row, col, cfa);
rgb[out_idx + 1] = green_at_rb_bilinear(data, width, height, row, col);
rgb[out_idx + 2] = val;
}
_ => {
rgb[out_idx + 1] = val;
}
}
}
}
rgb
}
#[cfg(feature = "rawloader")]
#[inline]
fn green_at_rb_bilinear(data: &[f32], width: usize, height: usize, row: usize, col: usize) -> f32 {
let mut sum = 0.0f32;
let mut count = 0u32;
if row > 0 {
sum += data[(row - 1) * width + col];
count += 1;
}
if row + 1 < height {
sum += data[(row + 1) * width + col];
count += 1;
}
if col > 0 {
sum += data[row * width + col - 1];
count += 1;
}
if col + 1 < width {
sum += data[row * width + col + 1];
count += 1;
}
if count > 0 { sum / count as f32 } else { 0.0 }
}
#[cfg(feature = "rawloader")]
#[inline]
fn rb_at_green_bilinear(
data: &[f32],
width: usize,
height: usize,
row: usize,
col: usize,
cfa: &rawloader::CFA,
) -> (f32, f32) {
let h_color = if col > 0 {
cfa.color_at(row, col - 1)
} else {
cfa.color_at(row, col + 1)
};
let h_avg = avg_horizontal(data, width, row, col);
let v_avg = avg_vertical(data, height, width, row, col);
if h_color == R {
(h_avg, v_avg)
} else {
(v_avg, h_avg)
}
}
#[cfg(feature = "rawloader")]
#[inline]
fn opposite_at_rb_bilinear(
data: &[f32],
width: usize,
height: usize,
row: usize,
col: usize,
cfa: &rawloader::CFA,
) -> f32 {
let _ = cfa; let mut sum = 0.0f32;
let mut count = 0u32;
if row > 0 && col > 0 {
sum += data[(row - 1) * width + col - 1];
count += 1;
}
if row > 0 && col + 1 < width {
sum += data[(row - 1) * width + col + 1];
count += 1;
}
if row + 1 < height && col > 0 {
sum += data[(row + 1) * width + col - 1];
count += 1;
}
if row + 1 < height && col + 1 < width {
sum += data[(row + 1) * width + col + 1];
count += 1;
}
if count > 0 { sum / count as f32 } else { 0.0 }
}
#[cfg(feature = "rawloader")]
#[inline]
fn avg_horizontal(data: &[f32], width: usize, row: usize, col: usize) -> f32 {
let left = if col > 0 {
data[row * width + col - 1]
} else {
0.0
};
let right = if col + 1 < width {
data[row * width + col + 1]
} else {
0.0
};
let count = (col > 0) as u32 + (col + 1 < width) as u32;
if count > 0 {
(left + right) / count as f32
} else {
0.0
}
}
#[cfg(feature = "rawloader")]
#[inline]
fn avg_vertical(data: &[f32], height: usize, width: usize, row: usize, col: usize) -> f32 {
let top = if row > 0 {
data[(row - 1) * width + col]
} else {
0.0
};
let bottom = if row + 1 < height {
data[(row + 1) * width + col]
} else {
0.0
};
let count = (row > 0) as u32 + (row + 1 < height) as u32;
if count > 0 {
(top + bottom) / count as f32
} else {
0.0
}
}
#[cfg(feature = "rawloader")]
fn demosaic_malvar(data: &[f32], width: usize, height: usize, cfa: &rawloader::CFA) -> Vec<f32> {
let mut rgb = vec![0.0f32; width * height * 3];
const BORDER: usize = 2;
if width <= 2 * BORDER || height <= 2 * BORDER {
for row in 0..height {
for col in 0..width {
malvar_pixel_clamped(data, &mut rgb, width, height, row, col, cfa);
}
}
return rgb;
}
let cfa_tile = [
[cfa.color_at(0, 0), cfa.color_at(0, 1)],
[cfa.color_at(1, 0), cfa.color_at(1, 1)],
];
let mut gh = [[0usize; 2]; 2];
for rp in 0..2 {
for cp in 0..2 {
if cfa_tile[rp][cp] == G {
gh[rp][cp] = cfa_tile[rp][1 - cp];
}
}
}
for row in 0..BORDER {
for col in 0..width {
malvar_pixel_clamped(data, &mut rgb, width, height, row, col, cfa);
}
}
for row in (height - BORDER)..height {
for col in 0..width {
malvar_pixel_clamped(data, &mut rgb, width, height, row, col, cfa);
}
}
for row in BORDER..(height - BORDER) {
for col in 0..BORDER {
malvar_pixel_clamped(data, &mut rgb, width, height, row, col, cfa);
}
for col in (width - BORDER)..width {
malvar_pixel_clamped(data, &mut rgb, width, height, row, col, cfa);
}
}
malvar_interior(data, &mut rgb, width, height, cfa_tile, gh);
rgb
}
#[cfg(feature = "rawloader")]
#[autoversion]
fn malvar_interior(
data: &[f32],
rgb: &mut [f32],
width: usize,
height: usize,
cfa_tile: [[usize; 2]; 2],
gh: [[usize; 2]; 2],
) {
const BORDER: usize = 2;
let w = width;
for row in BORDER..(height - BORDER) {
let rp = row & 1;
for col in BORDER..(width - BORDER) {
let cp = col & 1;
let color = cfa_tile[rp][cp];
let idx = row * w + col;
let out = idx * 3;
let val = data[idx];
let n = data[idx - w];
let s = data[idx + w];
let e = data[idx + 1];
let we = data[idx - 1];
let n2 = data[idx - 2 * w];
let s2 = data[idx + 2 * w];
let e2 = data[idx + 2];
let w2 = data[idx - 2];
let ne = data[idx - w + 1];
let nw = data[idx - w - 1];
let se = data[idx + w + 1];
let sw = data[idx + w - 1];
match color {
R => {
rgb[out] = val;
rgb[out + 1] =
((4.0 * val + 2.0 * (n + s + e + we) - (n2 + s2 + e2 + w2)) / 8.0).max(0.0);
rgb[out + 2] =
((6.0 * val + 2.0 * (ne + nw + se + sw) - 1.5 * (n2 + s2 + e2 + w2)) / 8.0)
.max(0.0);
}
G => {
let h = (5.0 * val + 4.0 * (e + we) - (ne + nw + se + sw) - (e2 + w2)
+ 0.5 * (n2 + s2))
/ 8.0;
let v = (5.0 * val + 4.0 * (n + s) - (ne + nw + se + sw) - (n2 + s2)
+ 0.5 * (e2 + w2))
/ 8.0;
rgb[out + 1] = val;
if gh[rp][cp] == R {
rgb[out] = h.max(0.0);
rgb[out + 2] = v.max(0.0);
} else {
rgb[out] = v.max(0.0);
rgb[out + 2] = h.max(0.0);
}
}
B => {
rgb[out] =
((6.0 * val + 2.0 * (ne + nw + se + sw) - 1.5 * (n2 + s2 + e2 + w2)) / 8.0)
.max(0.0);
rgb[out + 1] =
((4.0 * val + 2.0 * (n + s + e + we) - (n2 + s2 + e2 + w2)) / 8.0).max(0.0);
rgb[out + 2] = val;
}
_ => {
rgb[out + 1] = val;
}
}
}
}
}
#[cfg(feature = "rawloader")]
fn malvar_pixel_clamped(
data: &[f32],
rgb: &mut [f32],
width: usize,
height: usize,
row: usize,
col: usize,
cfa: &rawloader::CFA,
) {
let color = cfa.color_at(row, col);
let val = data[row * width + col];
let out_idx = (row * width + col) * 3;
match color {
R => {
rgb[out_idx] = val;
rgb[out_idx + 1] = malvar_green_at_rb(data, width, height, row, col);
rgb[out_idx + 2] = malvar_opposite_at_rb(data, width, height, row, col, cfa);
}
G => {
let (r, b) = malvar_rb_at_green(data, width, height, row, col, cfa);
rgb[out_idx] = r;
rgb[out_idx + 1] = val;
rgb[out_idx + 2] = b;
}
B => {
rgb[out_idx] = malvar_opposite_at_rb(data, width, height, row, col, cfa);
rgb[out_idx + 1] = malvar_green_at_rb(data, width, height, row, col);
rgb[out_idx + 2] = val;
}
_ => {
rgb[out_idx + 1] = val;
}
}
}
#[cfg(feature = "rawloader")]
#[inline]
fn px(data: &[f32], width: usize, height: usize, row: isize, col: isize) -> f32 {
let r = row.clamp(0, height as isize - 1) as usize;
let c = col.clamp(0, width as isize - 1) as usize;
data[r * width + c]
}
#[cfg(feature = "rawloader")]
#[inline]
fn malvar_green_at_rb(data: &[f32], width: usize, height: usize, row: usize, col: usize) -> f32 {
let r = row as isize;
let c = col as isize;
let center = px(data, width, height, r, c);
let n = px(data, width, height, r - 1, c);
let s = px(data, width, height, r + 1, c);
let e = px(data, width, height, r, c + 1);
let w = px(data, width, height, r, c - 1);
let n2 = px(data, width, height, r - 2, c);
let s2 = px(data, width, height, r + 2, c);
let e2 = px(data, width, height, r, c + 2);
let w2 = px(data, width, height, r, c - 2);
let val = (4.0 * center + 2.0 * (n + s + e + w) - (n2 + s2 + e2 + w2)) / 8.0;
val.max(0.0)
}
#[cfg(feature = "rawloader")]
#[inline]
fn malvar_opposite_at_rb(
data: &[f32],
width: usize,
height: usize,
row: usize,
col: usize,
cfa: &rawloader::CFA,
) -> f32 {
let _ = cfa;
let r = row as isize;
let c = col as isize;
let center = px(data, width, height, r, c);
let ne = px(data, width, height, r - 1, c + 1);
let nw = px(data, width, height, r - 1, c - 1);
let se = px(data, width, height, r + 1, c + 1);
let sw = px(data, width, height, r + 1, c - 1);
let n2 = px(data, width, height, r - 2, c);
let s2 = px(data, width, height, r + 2, c);
let e2 = px(data, width, height, r, c + 2);
let w2 = px(data, width, height, r, c - 2);
let val = (6.0 * center + 2.0 * (ne + nw + se + sw) - 1.5 * (n2 + s2 + e2 + w2)) / 8.0;
val.max(0.0)
}
#[cfg(feature = "rawloader")]
fn malvar_rb_at_green(
data: &[f32],
width: usize,
height: usize,
row: usize,
col: usize,
cfa: &rawloader::CFA,
) -> (f32, f32) {
let r = row as isize;
let c = col as isize;
let center = px(data, width, height, r, c);
let n = px(data, width, height, r - 1, c);
let s = px(data, width, height, r + 1, c);
let e = px(data, width, height, r, c + 1);
let w = px(data, width, height, r, c - 1);
let ne = px(data, width, height, r - 1, c + 1);
let nw = px(data, width, height, r - 1, c - 1);
let se = px(data, width, height, r + 1, c + 1);
let sw = px(data, width, height, r + 1, c - 1);
let n2 = px(data, width, height, r - 2, c);
let s2 = px(data, width, height, r + 2, c);
let e2 = px(data, width, height, r, c + 2);
let w2 = px(data, width, height, r, c - 2);
let h_val =
(5.0 * center + 4.0 * (e + w) - (ne + nw + se + sw) - (e2 + w2) + 0.5 * (n2 + s2)) / 8.0;
let v_val =
(5.0 * center + 4.0 * (n + s) - (ne + nw + se + sw) - (n2 + s2) + 0.5 * (e2 + w2)) / 8.0;
let h_color = if col > 0 {
cfa.color_at(row, col - 1)
} else {
cfa.color_at(row, col + 1)
};
if h_color == R {
(h_val.max(0.0), v_val.max(0.0))
} else {
(v_val.max(0.0), h_val.max(0.0))
}
}
#[cfg_attr(not(feature = "rawler"), allow(dead_code))]
pub(crate) struct CfaPattern {
colors: Vec<u8>,
width: usize,
height: usize,
}
#[cfg_attr(not(feature = "rawler"), allow(dead_code))]
impl CfaPattern {
pub fn new(colors: Vec<u8>, width: usize, height: usize) -> Self {
debug_assert_eq!(colors.len(), width * height);
Self {
colors,
width,
height,
}
}
#[inline]
pub fn color_at(&self, row: usize, col: usize) -> usize {
let r = row % self.height;
let c = col % self.width;
self.colors[r * self.width + c] as usize
}
}
#[cfg_attr(not(feature = "rawler"), allow(dead_code))]
pub(crate) fn demosaic_xtrans_bilinear(
data: &[f32],
width: usize,
height: usize,
cfa: &CfaPattern,
) -> Vec<f32> {
let mut rgb = vec![0.0f32; width * height * 3];
for row in 0..height {
for col in 0..width {
let known = cfa.color_at(row, col);
let out_idx = (row * width + col) * 3;
rgb[out_idx + known] = data[row * width + col];
for ch in 0..3 {
if ch == known {
continue;
}
rgb[out_idx + ch] =
interpolate_channel_xtrans(data, width, height, row, col, ch, cfa);
}
}
}
rgb
}
#[cfg_attr(not(feature = "rawler"), allow(dead_code))]
fn interpolate_channel_xtrans(
data: &[f32],
width: usize,
height: usize,
row: usize,
col: usize,
target_ch: usize,
cfa: &CfaPattern,
) -> f32 {
let mut sum = 0.0f32;
let mut count = 0u32;
let r_start = row.saturating_sub(2);
let r_end = (row + 3).min(height);
let c_start = col.saturating_sub(2);
let c_end = (col + 3).min(width);
for nr in r_start..r_end {
for nc in c_start..c_end {
if nr == row && nc == col {
continue;
}
if cfa.color_at(nr, nc) == target_ch {
sum += data[nr * width + nc];
count += 1;
}
}
}
if count > 0 { sum / count as f32 } else { 0.0 }
}
#[cfg(all(test, feature = "rawloader"))]
mod tests {
use super::*;
fn make_test_bayer() -> (Vec<f32>, usize, usize, rawloader::CFA) {
let cfa = rawloader::CFA::new("RGGB");
let width = 4;
let height = 4;
let mut data = vec![0.0f32; width * height];
for row in 0..height {
for col in 0..width {
let color = cfa.color_at(row, col);
data[row * width + col] = match color {
R => 0.8,
G => 0.5,
B => 0.3,
_ => 0.0,
};
}
}
(data, width, height, cfa)
}
#[test]
fn bilinear_produces_correct_dimensions() {
let (data, width, height, cfa) = make_test_bayer();
let rgb = demosaic_bilinear(&data, width, height, &cfa);
assert_eq!(rgb.len(), width * height * 3);
}
#[test]
fn malvar_produces_correct_dimensions() {
let (data, width, height, cfa) = make_test_bayer();
let rgb = demosaic_malvar(&data, width, height, &cfa);
assert_eq!(rgb.len(), width * height * 3);
}
#[test]
fn bilinear_known_channel_preserved() {
let (data, width, height, cfa) = make_test_bayer();
let rgb = demosaic_bilinear(&data, width, height, &cfa);
for row in 0..height {
for col in 0..width {
let color = cfa.color_at(row, col);
let idx = (row * width + col) * 3;
match color {
R => assert!((rgb[idx] - 0.8).abs() < 1e-6, "R at ({row},{col})"),
G => assert!((rgb[idx + 1] - 0.5).abs() < 1e-6, "G at ({row},{col})"),
B => assert!((rgb[idx + 2] - 0.3).abs() < 1e-6, "B at ({row},{col})"),
_ => {}
}
}
}
}
#[test]
fn malvar_known_channel_preserved() {
let (data, width, height, cfa) = make_test_bayer();
let rgb = demosaic_malvar(&data, width, height, &cfa);
for row in 0..height {
for col in 0..width {
let color = cfa.color_at(row, col);
let idx = (row * width + col) * 3;
match color {
R => assert!((rgb[idx] - 0.8).abs() < 1e-6, "R at ({row},{col})"),
G => assert!((rgb[idx + 1] - 0.5).abs() < 1e-6, "G at ({row},{col})"),
B => assert!((rgb[idx + 2] - 0.3).abs() < 1e-6, "B at ({row},{col})"),
_ => {}
}
}
}
}
#[test]
fn bilinear_output_non_negative() {
let (data, width, height, cfa) = make_test_bayer();
let rgb = demosaic_bilinear(&data, width, height, &cfa);
for val in &rgb {
assert!(*val >= 0.0, "Bilinear produced negative value: {val}");
}
}
#[test]
fn malvar_output_non_negative() {
let (data, width, height, cfa) = make_test_bayer();
let rgb = demosaic_malvar(&data, width, height, &cfa);
for val in &rgb {
assert!(*val >= 0.0, "Malvar produced negative value: {val}");
}
}
#[test]
fn uniform_input_produces_uniform_output() {
let cfa = rawloader::CFA::new("RGGB");
let width = 8;
let height = 8;
let data = vec![0.5f32; width * height];
let rgb_bilinear = demosaic_bilinear(&data, width, height, &cfa);
let rgb_malvar = demosaic_malvar(&data, width, height, &cfa);
for row in 2..height - 2 {
for col in 2..width - 2 {
let idx = (row * width + col) * 3;
for ch in 0..3 {
assert!(
(rgb_bilinear[idx + ch] - 0.5).abs() < 1e-4,
"Bilinear interior not uniform at ({row},{col}) ch={ch}: {}",
rgb_bilinear[idx + ch]
);
assert!(
(rgb_malvar[idx + ch] - 0.5).abs() < 1e-4,
"Malvar interior not uniform at ({row},{col}) ch={ch}: {}",
rgb_malvar[idx + ch]
);
}
}
}
}
#[test]
fn different_cfa_patterns() {
let width = 8;
let height = 8;
let data = vec![0.5f32; width * height];
for pattern in &["RGGB", "BGGR", "GRBG", "GBRG"] {
let cfa = rawloader::CFA::new(pattern);
let rgb = demosaic_to_rgb_f32(&data, width, height, &cfa, DemosaicMethod::Bilinear);
assert_eq!(rgb.len(), width * height * 3, "Pattern {pattern}");
let rgb =
demosaic_to_rgb_f32(&data, width, height, &cfa, DemosaicMethod::MalvarHeCutler);
assert_eq!(rgb.len(), width * height * 3, "Pattern {pattern}");
}
}
fn make_xtrans_cfa() -> CfaPattern {
#[rustfmt::skip]
let colors: Vec<u8> = vec![
G, B, G, G, R, G,
R, G, R, B, G, B,
G, B, G, G, R, G,
G, R, G, G, B, G,
B, G, B, R, G, R,
G, R, G, G, B, G,
].into_iter().map(|c| c as u8).collect();
CfaPattern::new(colors, 6, 6)
}
#[test]
fn xtrans_bilinear_dimensions() {
let cfa = make_xtrans_cfa();
let width = 24;
let height = 24;
let data = vec![0.5f32; width * height];
let rgb = demosaic_xtrans_bilinear(&data, width, height, &cfa);
assert_eq!(rgb.len(), width * height * 3);
}
#[test]
fn xtrans_bilinear_known_channel_preserved() {
let cfa = make_xtrans_cfa();
let width = 12;
let height = 12;
let mut data = vec![0.0f32; width * height];
for row in 0..height {
for col in 0..width {
data[row * width + col] = match cfa.color_at(row, col) {
R => 0.8,
G => 0.5,
B => 0.3,
_ => 0.0,
};
}
}
let rgb = demosaic_xtrans_bilinear(&data, width, height, &cfa);
for row in 0..height {
for col in 0..width {
let known = cfa.color_at(row, col);
let idx = (row * width + col) * 3;
let expected = match known {
R => 0.8,
G => 0.5,
B => 0.3,
_ => 0.0,
};
assert!(
(rgb[idx + known] - expected).abs() < 1e-6,
"Known channel not preserved at ({row},{col})"
);
}
}
}
#[test]
fn xtrans_bilinear_uniform_input() {
let cfa = make_xtrans_cfa();
let width = 24;
let height = 24;
let data = vec![0.5f32; width * height];
let rgb = demosaic_xtrans_bilinear(&data, width, height, &cfa);
for row in 3..height - 3 {
for col in 3..width - 3 {
let idx = (row * width + col) * 3;
for ch in 0..3 {
assert!(
(rgb[idx + ch] - 0.5).abs() < 1e-4,
"X-Trans interior not uniform at ({row},{col}) ch={ch}: {}",
rgb[idx + ch]
);
}
}
}
}
#[test]
fn xtrans_bilinear_non_negative() {
let cfa = make_xtrans_cfa();
let width = 12;
let height = 12;
let mut data = vec![0.0f32; width * height];
for (i, v) in data.iter_mut().enumerate() {
*v = ((i * 37 + 13) % 100) as f32 / 100.0;
}
let rgb = demosaic_xtrans_bilinear(&data, width, height, &cfa);
for val in &rgb {
assert!(*val >= 0.0, "X-Trans produced negative value: {val}");
}
}
}