Files
rustytorch/crates/specialized/rtx-registration/src/interpolate.rs
T
2026-03-04 00:08:42 +00:00

275 lines
7.7 KiB
Rust

//! Image interpolation for transformed coordinates.
use nalgebra::Point3;
use rtx_medical_io::Volume;
/// Interpolation method.
#[derive(Debug, Clone, Copy, Default)]
pub enum InterpolationMethod {
/// Nearest neighbor interpolation.
NearestNeighbor,
/// Trilinear interpolation.
#[default]
Trilinear,
}
/// Interpolate a volume at a given physical coordinate.
pub fn interpolate(volume: &Volume, point: &Point3<f64>, method: InterpolationMethod) -> f64 {
match method {
InterpolationMethod::NearestNeighbor => interpolate_nearest(volume, point),
InterpolationMethod::Trilinear => interpolate_trilinear(volume, point),
}
}
/// Nearest neighbor interpolation.
pub fn interpolate_nearest(volume: &Volume, point: &Point3<f64>) -> f64 {
let shape = volume.shape();
let spacing = volume.spacing();
// Convert physical to voxel coordinates
let x = (point.x / spacing[0]).round() as i64;
let y = (point.y / spacing[1]).round() as i64;
let z = (point.z / spacing[2]).round() as i64;
// Bounds check
if x < 0
|| x >= shape[0] as i64
|| y < 0
|| y >= shape[1] as i64
|| z < 0
|| z >= shape[2] as i64
{
return 0.0;
}
volume
.get(x as usize, y as usize, z as usize)
.unwrap_or(0.0)
}
/// Trilinear interpolation.
pub fn interpolate_trilinear(volume: &Volume, point: &Point3<f64>) -> f64 {
let shape = volume.shape();
let spacing = volume.spacing();
// Convert physical to voxel coordinates
let x = point.x / spacing[0];
let y = point.y / spacing[1];
let z = point.z / spacing[2];
// Get integer and fractional parts
let x0 = x.floor() as i64;
let y0 = y.floor() as i64;
let z0 = z.floor() as i64;
let x1 = x0 + 1;
let y1 = y0 + 1;
let z1 = z0 + 1;
let xd = x - x0 as f64;
let yd = y - y0 as f64;
let zd = z - z0 as f64;
// Bounds check
if x0 < 0
|| x1 >= shape[0] as i64
|| y0 < 0
|| y1 >= shape[1] as i64
|| z0 < 0
|| z1 >= shape[2] as i64
{
return 0.0;
}
let x0 = x0 as usize;
let y0 = y0 as usize;
let z0 = z0 as usize;
let x1 = x1 as usize;
let y1 = y1 as usize;
let z1 = z1 as usize;
// Get corner values
let c000 = volume.get(x0, y0, z0).unwrap_or(0.0);
let c001 = volume.get(x0, y0, z1).unwrap_or(0.0);
let c010 = volume.get(x0, y1, z0).unwrap_or(0.0);
let c011 = volume.get(x0, y1, z1).unwrap_or(0.0);
let c100 = volume.get(x1, y0, z0).unwrap_or(0.0);
let c101 = volume.get(x1, y0, z1).unwrap_or(0.0);
let c110 = volume.get(x1, y1, z0).unwrap_or(0.0);
let c111 = volume.get(x1, y1, z1).unwrap_or(0.0);
// Interpolate along x
let c00 = c000 * (1.0 - xd) + c100 * xd;
let c01 = c001 * (1.0 - xd) + c101 * xd;
let c10 = c010 * (1.0 - xd) + c110 * xd;
let c11 = c011 * (1.0 - xd) + c111 * xd;
// Interpolate along y
let c0 = c00 * (1.0 - yd) + c10 * yd;
let c1 = c01 * (1.0 - yd) + c11 * yd;
// Interpolate along z
c0 * (1.0 - zd) + c1 * zd
}
/// Apply a transform and resample a volume.
pub fn resample_volume<T: crate::transform::Transform>(
volume: &Volume,
transform: &T,
output_shape: [usize; 3],
output_spacing: [f64; 3],
method: InterpolationMethod,
) -> Volume {
let mut output = Volume::zeros(output_shape);
output.set_spacing(output_spacing);
for z in 0..output_shape[2] {
for y in 0..output_shape[1] {
for x in 0..output_shape[0] {
// Physical coordinate in output space
let point = Point3::new(
x as f64 * output_spacing[0],
y as f64 * output_spacing[1],
z as f64 * output_spacing[2],
);
// Transform to input space (use inverse transform)
if let Some(inv) = transform.try_inverse() {
let input_point = inv.transform_point(&point);
let value = interpolate(volume, &input_point, method);
output.set(x, y, z, value);
}
}
}
}
output
}
/// Apply a transform and resample a volume in parallel.
pub fn resample_volume_parallel<T: crate::transform::Transform>(
volume: &Volume,
transform: &T,
output_shape: [usize; 3],
output_spacing: [f64; 3],
method: InterpolationMethod,
) -> Volume {
use rayon::prelude::*;
let inv = match transform.try_inverse() {
Some(inv) => inv,
None => return Volume::zeros(output_shape),
};
let total_voxels = output_shape[0] * output_shape[1] * output_shape[2];
let values: Vec<f64> = (0..total_voxels)
.into_par_iter()
.map(|idx| {
let x = idx % output_shape[0];
let y = (idx / output_shape[0]) % output_shape[1];
let z = idx / (output_shape[0] * output_shape[1]);
let point = Point3::new(
x as f64 * output_spacing[0],
y as f64 * output_spacing[1],
z as f64 * output_spacing[2],
);
let input_point = inv.transform_point(&point);
interpolate(volume, &input_point, method)
})
.collect();
let mut output = Volume::zeros(output_shape);
output.set_spacing(output_spacing);
for (idx, value) in values.into_iter().enumerate() {
let x = idx % output_shape[0];
let y = (idx / output_shape[0]) % output_shape[1];
let z = idx / (output_shape[0] * output_shape[1]);
output.set(x, y, z, value);
}
output
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_nearest_neighbor() {
let mut vol = Volume::zeros([10, 10, 10]);
vol.set(5, 5, 5, 100.0);
// Exact voxel center
let p = Point3::new(5.0, 5.0, 5.0);
let v = interpolate_nearest(&vol, &p);
assert_eq!(v, 100.0);
// Slightly off should still round to (5,5,5)
let p = Point3::new(5.2, 4.8, 5.3);
let v = interpolate_nearest(&vol, &p);
assert_eq!(v, 100.0);
// Far off should be 0
let p = Point3::new(0.0, 0.0, 0.0);
let v = interpolate_nearest(&vol, &p);
assert_eq!(v, 0.0);
}
#[test]
fn test_trilinear_center() {
let mut vol = Volume::zeros([10, 10, 10]);
vol.set(5, 5, 5, 100.0);
// Exact voxel position
let p = Point3::new(5.0, 5.0, 5.0);
let v = interpolate_trilinear(&vol, &p);
assert_eq!(v, 100.0);
}
#[test]
fn test_trilinear_midpoint() {
let mut vol = Volume::zeros([10, 10, 10]);
// Set two adjacent voxels
vol.set(5, 5, 5, 0.0);
vol.set(6, 5, 5, 100.0);
// Midpoint along x
let p = Point3::new(5.5, 5.0, 5.0);
let v = interpolate_trilinear(&vol, &p);
assert!((v - 50.0).abs() < 1e-10);
}
#[test]
fn test_trilinear_corner() {
let mut vol = Volume::zeros([10, 10, 10]);
// Set all 8 corners of a cube
for dz in 0..2 {
for dy in 0..2 {
for dx in 0..2 {
vol.set(2 + dx, 2 + dy, 2 + dz, (dx + dy + dz) as f64 * 100.0);
}
}
}
// Center of cube
let p = Point3::new(2.5, 2.5, 2.5);
let v = interpolate_trilinear(&vol, &p);
// Average of 0, 100, 100, 200, 100, 200, 200, 300 = 150
assert!((v - 150.0).abs() < 1e-10);
}
#[test]
fn test_out_of_bounds() {
let vol = Volume::zeros([10, 10, 10]);
let p = Point3::new(-1.0, 5.0, 5.0);
assert_eq!(interpolate_trilinear(&vol, &p), 0.0);
let p = Point3::new(10.0, 5.0, 5.0);
assert_eq!(interpolate_trilinear(&vol, &p), 0.0);
}
}