use std::io::{BufWriter, Write};
use fft2d::nalgebra::dcst::{dct_2d, idct_2d};
use image::{GrayImage, ImageBuffer, Luma, Primitive, Rgb};
use nalgebra::{
allocator::Allocator, DMatrix, DefaultAllocator, Dim, MatrixSlice, Scalar, Vector2, Vector3,
};
use show_image::create_window;
#[show_image::main]
fn main() -> Result<(), Box<dyn std::error::Error>> {
let img = image::open("data/cat-normal-map.png")?.into_rgb8();
let window_in = create_window("input normals", Default::default())?;
window_in.set_image("Input image", img.clone())?;
window_in.wait_until_destroyed()?;
let mut normals = matrix_from_rgb_image(&img, |x| *x as f32 / 255.0);
for n in normals.iter_mut() {
if n.x + n.y + n.z != 0.0 {
n.x = (n.x - 0.5) * -2.0;
n.y = (n.y - 0.5) * -2.0;
n.z = (n.z - 0.5).max(0.001) * -2.0;
n.normalize_mut();
}
}
let depths = normal_integration(&normals);
mat_show("depth", depths.slice_range(.., ..));
eprintln!("Saving data/cat.obj to disk");
save_as_obj("data/cat.obj", &depths).unwrap();
let depth_min = depths.min();
let depth_max = depths.max();
eprintln!("depths within [ {}, {} ]", depth_min, depth_max);
let depths_to_gray =
|z| ((z - depth_min) / (depth_max - depth_min) * (256.0 * 256.0 - 1.0)) as u16;
let depth_img = image_from_matrix(&depths, depths_to_gray);
eprintln!("Saving data/cat-depth.png to disk");
depth_img.save("data/cat-depth.png").unwrap();
Ok(())
}
fn normal_integration(normals: &DMatrix<Vector3<f32>>) -> DMatrix<f32> {
let (nrows, ncols) = normals.shape();
let mut gradient_x = DMatrix::zeros(nrows, ncols);
let mut gradient_y = DMatrix::zeros(nrows, ncols);
for ((gx, gy), n) in gradient_x
.iter_mut()
.zip(gradient_y.iter_mut())
.zip(normals.iter())
{
if n.z < -0.01 {
*gx = -n.x / n.z;
*gy = -n.y / n.z;
}
}
mat_show("gx", gradient_x.slice_range(.., ..));
mat_show("gy", gradient_y.slice_range(.., ..));
dct_poisson(&gradient_y, &gradient_x)
}
fn dct_poisson(p: &DMatrix<f32>, q: &DMatrix<f32>) -> DMatrix<f32> {
let (nrows, ncols) = p.shape();
let mut px = DMatrix::zeros(nrows, ncols);
let p_top = p.slice_range(0..nrows - 2, ..);
let p_bottom = p.slice_range(2..nrows, ..);
px.slice_range_mut(1..nrows - 1, ..)
.copy_from(&(0.5 * (p_bottom - p_top)));
px.row_mut(0).copy_from(&(0.5 * (p.row(1) - p.row(0))));
px.row_mut(nrows - 1)
.copy_from(&(0.5 * (p.row(nrows - 1) - p.row(nrows - 2))));
let mut qy = DMatrix::zeros(nrows, ncols);
let q_left = q.slice_range(.., 0..ncols - 2);
let q_right = q.slice_range(.., 2..ncols);
qy.slice_range_mut(.., 1..ncols - 1)
.copy_from(&(0.5 * (q_right - q_left)));
qy.column_mut(0)
.copy_from(&(0.5 * (q.column(1) - q.column(0))));
qy.column_mut(ncols - 1)
.copy_from(&(0.5 * (q.column(ncols - 1) - q.column(ncols - 2))));
let mut f = px + qy;
let mut f_left = f.column_mut(0);
f_left += q.column(0);
let mut f_top = f.row_mut(0);
f_top += p.row(0);
let mut f_bottom = f.row_mut(nrows - 1);
f_bottom -= p.row(nrows - 1);
let mut f_right = f.column_mut(ncols - 1);
f_right -= q.column(ncols - 1);
let mut f_cos = dct_2d(f.map(|x| x as f64));
let pi = std::f64::consts::PI;
let coords = coordinates_column_major((ncols, nrows));
for (f_cos_ij, (i, j)) in f_cos.iter_mut().zip(coords) {
let v = Vector2::new(i as f64 / ncols as f64, j as f64 / nrows as f64);
let denom = 4.0 * v.map(|u| (0.5 * pi * u).sin()).norm_squared();
*f_cos_ij /= -(denom.max(1e-7));
}
let depths = idct_2d(f_cos);
let z_min = depths.min();
let dct_norm_coef = 4.0 / (nrows * ncols) as f32;
depths.map(|z| dct_norm_coef * (z - z_min) as f32)
}
fn coordinates_column_major(shape: (usize, usize)) -> impl Iterator<Item = (usize, usize)> {
let (nrows, ncols) = shape;
(0..ncols)
.map(move |j| (0..nrows).map(move |i| (i, j)))
.flatten()
}
fn mat_show<'a, R, C, RStride, CStride>(
title: &str,
mat: MatrixSlice<'a, f32, R, C, RStride, CStride>,
) where
R: Dim,
C: Dim,
RStride: Dim,
CStride: Dim,
DefaultAllocator: Allocator<f32, C, R>,
{
let mat = mat.to_owned().transpose();
let mat_min = mat.min();
let mat_max = mat.max();
let mat_mean = mat.mean();
eprintln!(
"{} values: min: {}, mean: {}, max: {}",
title, mat_min, mat_mean, mat_max
);
let img_data: Vec<u8> = mat
.iter()
.map(|x| ((x - mat_min) / (mat_max - mat_min) * 255.0) as u8)
.collect();
let (width, height) = mat.shape();
let img_transposed = GrayImage::from_raw(width as u32, height as u32, img_data).unwrap();
let window = create_window(title, Default::default()).unwrap();
window.set_image("", img_transposed).unwrap();
window.wait_until_destroyed().unwrap();
}
fn save_as_obj(path: &str, depth_map: &DMatrix<f32>) -> std::io::Result<()> {
let z_min = depth_map.min();
let z_max = depth_map.max();
let transform = |z| z_min + z_max - z;
let file = std::fs::File::create(path)?;
let mut buf_writer = BufWriter::new(file);
let (height, width) = depth_map.shape();
let scale = z_max - z_min;
let precision = (-(scale / 65536.0).log10()).ceil().max(0.0) as usize;
for ((y, x), z) in coordinates_column_major((height, width)).zip(depth_map) {
writeln!(
&mut buf_writer,
"v {} -{} {:.prec$}",
x,
y,
transform(z),
prec = precision,
)?;
}
for (y, x) in coordinates_column_major((height - 1, width - 1)) {
let top_left = height * x + y + 1; let bot_left = top_left + 1;
let bot_right = bot_left + height;
let top_right = bot_right - 1;
writeln!(
&mut buf_writer,
"f {} {} {} {}",
top_left, bot_left, bot_right, top_right,
)?;
}
Ok(())
}
#[allow(clippy::cast_possible_truncation)]
fn image_from_matrix<'a, T: Scalar, U: 'static + Primitive, F: Fn(&'a T) -> U>(
mat: &'a DMatrix<T>,
to_gray: F,
) -> ImageBuffer<Luma<U>, Vec<U>> {
let (nb_rows, nb_cols) = mat.shape();
let mut img_buf = ImageBuffer::new(nb_cols as u32, nb_rows as u32);
for (x, y, pixel) in img_buf.enumerate_pixels_mut() {
*pixel = Luma([to_gray(&mat[(y as usize, x as usize)])]);
}
img_buf
}
fn matrix_from_rgb_image<'a, T: 'static + Primitive, U: Scalar, F: Fn(&'a T) -> U>(
img: &'a ImageBuffer<Rgb<T>, Vec<T>>,
scale: F,
) -> DMatrix<Vector3<U>> {
let (width, height) = img.dimensions();
DMatrix::from_iterator(
width as usize,
height as usize,
img.as_raw()
.chunks_exact(3)
.map(|s| Vector3::new(scale(&s[0]), scale(&s[1]), scale(&s[2]))),
)
.transpose()
}