use crate::error::{NdimageError, NdimageResult};
use scirs2_core::ndarray::{Array1, Array3};
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Padding3D {
Zero,
Reflect,
Replicate,
}
#[inline]
fn reflect_index(i: isize, n: usize) -> Option<usize> {
if n == 0 {
return None;
}
let n = n as isize;
let period = 2 * n - 2;
if period <= 0 {
return Some(0); }
let mut r = i % period;
if r < 0 {
r += period;
}
if r >= n {
r = period - r;
}
Some(r as usize)
}
#[inline]
fn resolve_index(i: isize, n: usize, padding: Padding3D) -> Option<usize> {
if i >= 0 && (i as usize) < n {
return Some(i as usize);
}
match padding {
Padding3D::Zero => None,
Padding3D::Replicate => {
if n == 0 {
None
} else {
Some(i.clamp(0, n as isize - 1) as usize)
}
}
Padding3D::Reflect => reflect_index(i, n),
}
}
#[inline]
fn get_padded(volume: &Array3<f64>, z: isize, y: isize, x: isize, padding: Padding3D) -> f64 {
let shape = volume.shape();
let oz = resolve_index(z, shape[0], padding);
let oy = resolve_index(y, shape[1], padding);
let ox = resolve_index(x, shape[2], padding);
match (oz, oy, ox) {
(Some(rz), Some(ry), Some(rx)) => volume[[rz, ry, rx]],
_ => 0.0, }
}
pub fn convolve3d(
volume: &Array3<f64>,
kernel: &Array3<f64>,
padding: Padding3D,
) -> NdimageResult<Array3<f64>> {
let ks = kernel.shape();
if ks[0] == 0 || ks[1] == 0 || ks[2] == 0 {
return Err(NdimageError::InvalidInput(
"kernel must have non-zero extent on all axes".to_string(),
));
}
let vs = volume.shape();
let (sz, sy, sx) = (vs[0], vs[1], vs[2]);
let (kz, ky, kx) = (ks[0], ks[1], ks[2]);
let hz = (kz / 2) as isize;
let hy = (ky / 2) as isize;
let hx = (kx / 2) as isize;
let mut out = Array3::zeros((sz, sy, sx));
for iz in 0..sz {
for iy in 0..sy {
for ix in 0..sx {
let mut acc = 0.0f64;
for dkz in 0..kz {
for dky in 0..ky {
for dkx in 0..kx {
let fkz = kz - 1 - dkz;
let fky = ky - 1 - dky;
let fkx = kx - 1 - dkx;
let kval = kernel[[fkz, fky, fkx]];
let vz = iz as isize + dkz as isize - hz;
let vy = iy as isize + dky as isize - hy;
let vx = ix as isize + dkx as isize - hx;
acc += kval * get_padded(volume, vz, vy, vx, padding);
}
}
}
out[[iz, iy, ix]] = acc;
}
}
}
Ok(out)
}
pub fn correlate3d(
volume: &Array3<f64>,
kernel: &Array3<f64>,
padding: Padding3D,
) -> NdimageResult<Array3<f64>> {
let ks = kernel.shape();
if ks[0] == 0 || ks[1] == 0 || ks[2] == 0 {
return Err(NdimageError::InvalidInput(
"kernel must have non-zero extent on all axes".to_string(),
));
}
let vs = volume.shape();
let (sz, sy, sx) = (vs[0], vs[1], vs[2]);
let (kz, ky, kx) = (ks[0], ks[1], ks[2]);
let hz = (kz / 2) as isize;
let hy = (ky / 2) as isize;
let hx = (kx / 2) as isize;
let mut out = Array3::zeros((sz, sy, sx));
for iz in 0..sz {
for iy in 0..sy {
for ix in 0..sx {
let mut acc = 0.0f64;
for dkz in 0..kz {
for dky in 0..ky {
for dkx in 0..kx {
let kval = kernel[[dkz, dky, dkx]];
let vz = iz as isize + dkz as isize - hz;
let vy = iy as isize + dky as isize - hy;
let vx = ix as isize + dkx as isize - hx;
acc += kval * get_padded(volume, vz, vy, vx, padding);
}
}
}
out[[iz, iy, ix]] = acc;
}
}
}
Ok(out)
}
fn convolve1d_z(src: &Array3<f64>, kernel: &Array1<f64>, padding: Padding3D) -> Array3<f64> {
let (sz, sy, sx) = (src.shape()[0], src.shape()[1], src.shape()[2]);
let klen = kernel.len();
let half = (klen / 2) as isize;
let mut out = Array3::zeros((sz, sy, sx));
for iz in 0..sz {
for iy in 0..sy {
for ix in 0..sx {
let mut acc = 0.0f64;
for (ki, &kv) in kernel.iter().enumerate() {
let nz = iz as isize + ki as isize - half;
acc += kv * get_padded(src, nz, iy as isize, ix as isize, padding);
}
out[[iz, iy, ix]] = acc;
}
}
}
out
}
fn convolve1d_y(src: &Array3<f64>, kernel: &Array1<f64>, padding: Padding3D) -> Array3<f64> {
let (sz, sy, sx) = (src.shape()[0], src.shape()[1], src.shape()[2]);
let klen = kernel.len();
let half = (klen / 2) as isize;
let mut out = Array3::zeros((sz, sy, sx));
for iz in 0..sz {
for iy in 0..sy {
for ix in 0..sx {
let mut acc = 0.0f64;
for (ki, &kv) in kernel.iter().enumerate() {
let ny = iy as isize + ki as isize - half;
acc += kv * get_padded(src, iz as isize, ny, ix as isize, padding);
}
out[[iz, iy, ix]] = acc;
}
}
}
out
}
fn convolve1d_x(src: &Array3<f64>, kernel: &Array1<f64>, padding: Padding3D) -> Array3<f64> {
let (sz, sy, sx) = (src.shape()[0], src.shape()[1], src.shape()[2]);
let klen = kernel.len();
let half = (klen / 2) as isize;
let mut out = Array3::zeros((sz, sy, sx));
for iz in 0..sz {
for iy in 0..sy {
for ix in 0..sx {
let mut acc = 0.0f64;
for (ki, &kv) in kernel.iter().enumerate() {
let nx = ix as isize + ki as isize - half;
acc += kv * get_padded(src, iz as isize, iy as isize, nx, padding);
}
out[[iz, iy, ix]] = acc;
}
}
}
out
}
fn gaussian_kernel1d(sigma: f64, truncate: f64) -> Array1<f64> {
let radius = (truncate * sigma).ceil() as usize;
let size = 2 * radius + 1;
let mut k = Array1::zeros(size);
let two_s2 = 2.0 * sigma * sigma;
let mut sum = 0.0f64;
for i in 0..size {
let x = i as f64 - radius as f64;
let v = (-x * x / two_s2).exp();
k[i] = v;
sum += v;
}
if sum != 0.0 {
k.mapv_inplace(|v| v / sum);
}
k
}
pub fn gaussian_filter3d(
volume: &Array3<f64>,
sigma: [f64; 3],
padding: Padding3D,
) -> NdimageResult<Array3<f64>> {
for (axis, &s) in sigma.iter().enumerate() {
if s <= 0.0 {
return Err(NdimageError::InvalidInput(format!(
"sigma[{}] must be positive, got {}",
axis, s
)));
}
}
const TRUNCATE: f64 = 4.0;
let kz = gaussian_kernel1d(sigma[0], TRUNCATE);
let ky = gaussian_kernel1d(sigma[1], TRUNCATE);
let kx = gaussian_kernel1d(sigma[2], TRUNCATE);
let tmp1 = convolve1d_z(volume, &kz, padding);
let tmp2 = convolve1d_y(&tmp1, &ky, padding);
Ok(convolve1d_x(&tmp2, &kx, padding))
}
pub fn uniform_filter3d(
volume: &Array3<f64>,
size: [usize; 3],
padding: Padding3D,
) -> NdimageResult<Array3<f64>> {
for (axis, &s) in size.iter().enumerate() {
if s == 0 {
return Err(NdimageError::InvalidInput(format!(
"size[{}] must be at least 1",
axis
)));
}
}
let make_box = |s: usize| -> Array1<f64> {
Array1::from_elem(s, 1.0 / s as f64)
};
let kz = make_box(size[0]);
let ky = make_box(size[1]);
let kx = make_box(size[2]);
let tmp1 = convolve1d_z(volume, &kz, padding);
let tmp2 = convolve1d_y(&tmp1, &ky, padding);
Ok(convolve1d_x(&tmp2, &kx, padding))
}
pub fn median_filter3d(
volume: &Array3<f64>,
size: [usize; 3],
padding: Padding3D,
) -> NdimageResult<Array3<f64>> {
for (axis, &s) in size.iter().enumerate() {
if s == 0 {
return Err(NdimageError::InvalidInput(format!(
"size[{}] must be at least 1",
axis
)));
}
}
let vs = volume.shape();
let (sz, sy, sx) = (vs[0], vs[1], vs[2]);
let hz = (size[0] / 2) as isize;
let hy = (size[1] / 2) as isize;
let hx = (size[2] / 2) as isize;
let window_vol = size[0] * size[1] * size[2];
let mut out = Array3::zeros((sz, sy, sx));
for iz in 0..sz {
for iy in 0..sy {
for ix in 0..sx {
let mut window: Vec<f64> = Vec::with_capacity(window_vol);
for dz in -hz..=hz {
for dy in -hy..=hy {
for dx in -hx..=hx {
window.push(get_padded(
volume,
iz as isize + dz,
iy as isize + dy,
ix as isize + dx,
padding,
));
}
}
}
window.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));
out[[iz, iy, ix]] = window[window.len() / 2];
}
}
}
Ok(out)
}
pub fn laplacian_of_gaussian3d(
volume: &Array3<f64>,
sigma: [f64; 3],
padding: Padding3D,
) -> NdimageResult<Array3<f64>> {
let smoothed = gaussian_filter3d(volume, sigma, padding)?;
laplacian3d_internal(&smoothed, padding)
}
fn laplacian3d_internal(volume: &Array3<f64>, padding: Padding3D) -> NdimageResult<Array3<f64>> {
let vs = volume.shape();
let (sz, sy, sx) = (vs[0], vs[1], vs[2]);
let mut out = Array3::zeros((sz, sy, sx));
for iz in 0..sz {
for iy in 0..sy {
for ix in 0..sx {
let z = iz as isize;
let y = iy as isize;
let x = ix as isize;
let center = volume[[iz, iy, ix]];
let lp = get_padded(volume, z - 1, y, x, padding)
+ get_padded(volume, z + 1, y, x, padding)
+ get_padded(volume, z, y - 1, x, padding)
+ get_padded(volume, z, y + 1, x, padding)
+ get_padded(volume, z, y, x - 1, padding)
+ get_padded(volume, z, y, x + 1, padding)
- 6.0 * center;
out[[iz, iy, ix]] = lp;
}
}
}
Ok(out)
}
#[cfg(test)]
mod tests {
use super::*;
use scirs2_core::ndarray::Array3;
fn make_impulse(sz: usize, sy: usize, sx: usize) -> Array3<f64> {
let mut v = Array3::zeros((sz, sy, sx));
v[[sz / 2, sy / 2, sx / 2]] = 1.0;
v
}
#[test]
fn test_convolve3d_identity_kernel() {
let vol = Array3::<f64>::from_shape_fn((4, 5, 6), |(z, y, x)| {
(z * 30 + y * 6 + x) as f64
});
let mut identity = Array3::zeros((3, 3, 3));
identity[[1, 1, 1]] = 1.0;
let out = convolve3d(&vol, &identity, Padding3D::Zero).expect("convolve3d failed");
for iz in 1..3usize {
for iy in 1..4usize {
for ix in 1..5usize {
assert!(
(out[[iz, iy, ix]] - vol[[iz, iy, ix]]).abs() < 1e-12,
"mismatch at [{}, {}, {}]: got {}, expected {}",
iz, iy, ix, out[[iz, iy, ix]], vol[[iz, iy, ix]]
);
}
}
}
}
#[test]
fn test_correlate3d_vs_convolve3d_symmetric() {
let vol = Array3::<f64>::ones((6, 6, 6));
let kernel = Array3::from_shape_fn((3, 3, 3), |(z, y, x)| {
let dz = z as f64 - 1.0;
let dy = y as f64 - 1.0;
let dx = x as f64 - 1.0;
1.0 / (1.0 + dz * dz + dy * dy + dx * dx)
});
let c1 = convolve3d(&vol, &kernel, Padding3D::Replicate).expect("convolve3d");
let c2 = correlate3d(&vol, &kernel, Padding3D::Replicate).expect("correlate3d");
for iz in 0..6 {
for iy in 0..6 {
for ix in 0..6 {
assert!(
(c1[[iz, iy, ix]] - c2[[iz, iy, ix]]).abs() < 1e-12,
"asymmetry at [{},{},{}]",
iz, iy, ix
);
}
}
}
}
#[test]
fn test_gaussian_filter3d_uniform_volume() {
let vol = Array3::<f64>::ones((10, 10, 10));
let out = gaussian_filter3d(&vol, [1.0, 1.0, 1.0], Padding3D::Reflect)
.expect("gaussian_filter3d");
for iz in 3..7 {
for iy in 3..7 {
for ix in 3..7 {
assert!(
(out[[iz, iy, ix]] - 1.0).abs() < 1e-10,
"interior mismatch at [{},{},{}]: {}",
iz, iy, ix, out[[iz, iy, ix]]
);
}
}
}
}
#[test]
fn test_gaussian_filter3d_invalid_sigma() {
let vol = Array3::<f64>::ones((4, 4, 4));
assert!(gaussian_filter3d(&vol, [0.0, 1.0, 1.0], Padding3D::Zero).is_err());
assert!(gaussian_filter3d(&vol, [1.0, -1.0, 1.0], Padding3D::Zero).is_err());
}
#[test]
fn test_uniform_filter3d_mean() {
let vol = Array3::<f64>::ones((8, 8, 8));
let out = uniform_filter3d(&vol, [3, 3, 3], Padding3D::Replicate)
.expect("uniform_filter3d");
for iz in 1..7 {
for iy in 1..7 {
for ix in 1..7 {
assert!((out[[iz, iy, ix]] - 1.0).abs() < 1e-12);
}
}
}
}
#[test]
fn test_median_filter3d_constant() {
let vol = Array3::<f64>::from_elem((5, 5, 5), 7.0);
let out = median_filter3d(&vol, [3, 3, 3], Padding3D::Replicate)
.expect("median_filter3d");
for &v in out.iter() {
assert!((v - 7.0).abs() < 1e-12);
}
}
#[test]
fn test_median_filter3d_invalid_size() {
let vol = Array3::<f64>::ones((4, 4, 4));
assert!(median_filter3d(&vol, [0, 3, 3], Padding3D::Zero).is_err());
}
#[test]
fn test_log3d_impulse() {
let vol = make_impulse(12, 12, 12);
let log = laplacian_of_gaussian3d(&vol, [1.0, 1.0, 1.0], Padding3D::Zero)
.expect("laplacian_of_gaussian3d");
let center = log[[6, 6, 6]];
let neighbour = log[[7, 6, 6]];
assert!(center < neighbour, "center={}, neighbour={}", center, neighbour);
}
#[test]
fn test_padding_modes_shape() {
let vol = Array3::<f64>::ones((4, 5, 6));
let k = Array3::<f64>::ones((3, 3, 3));
for &mode in &[Padding3D::Zero, Padding3D::Reflect, Padding3D::Replicate] {
let out = convolve3d(&vol, &k, mode).expect("convolve3d");
assert_eq!(out.shape(), [4, 5, 6]);
}
}
#[test]
fn test_reflect_index_symmetry() {
assert_eq!(reflect_index(-1, 5), Some(1));
assert_eq!(reflect_index(5, 5), Some(3));
assert_eq!(reflect_index(0, 5), Some(0));
assert_eq!(reflect_index(4, 5), Some(4));
}
}