496 lines
15 KiB
Rust
496 lines
15 KiB
Rust
//! 3D volumetric data structure for medical imaging.
|
|
//!
|
|
//! The `Volume` struct provides a container for 3D image data with spatial
|
|
//! metadata (spacing, origin, affine transform).
|
|
|
|
use crate::nifti::transform::{self, Affine4};
|
|
use rayon::prelude::*;
|
|
use serde::{Deserialize, Serialize};
|
|
|
|
/// A 3D volume of medical imaging data.
|
|
///
|
|
/// Data is stored in row-major order (C-order): z varies slowest, x varies fastest.
|
|
/// The layout is `data[z * (y_size * x_size) + y * x_size + x]`.
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct Volume {
|
|
/// Voxel data in row-major order
|
|
data: Vec<f64>,
|
|
|
|
/// Dimensions [x, y, z]
|
|
shape: [usize; 3],
|
|
|
|
/// Voxel spacing in mm [dx, dy, dz]
|
|
spacing: [f64; 3],
|
|
|
|
/// Origin in world coordinates [ox, oy, oz]
|
|
origin: [f64; 3],
|
|
|
|
/// 4x4 affine transformation matrix (voxel to world)
|
|
affine: Affine4,
|
|
}
|
|
|
|
impl Volume {
|
|
/// Create a new volume from raw data and metadata.
|
|
///
|
|
/// # Arguments
|
|
/// * `data` - Voxel data in row-major order
|
|
/// * `shape` - Dimensions [x, y, z]
|
|
/// * `spacing` - Voxel spacing in mm [dx, dy, dz]
|
|
/// * `origin` - Origin in world coordinates [ox, oy, oz]
|
|
/// * `affine` - 4x4 affine transformation matrix
|
|
///
|
|
/// # Panics
|
|
/// Panics if data length doesn't match shape.
|
|
pub fn new(
|
|
data: Vec<f64>,
|
|
shape: [usize; 3],
|
|
spacing: [f64; 3],
|
|
origin: [f64; 3],
|
|
affine: Affine4,
|
|
) -> Self {
|
|
let expected_len = shape[0] * shape[1] * shape[2];
|
|
assert_eq!(
|
|
data.len(),
|
|
expected_len,
|
|
"Data length {} doesn't match shape {:?} (expected {})",
|
|
data.len(),
|
|
shape,
|
|
expected_len
|
|
);
|
|
|
|
Self {
|
|
data,
|
|
shape,
|
|
spacing,
|
|
origin,
|
|
affine,
|
|
}
|
|
}
|
|
|
|
/// Create a volume filled with zeros.
|
|
pub fn zeros(shape: [usize; 3]) -> Self {
|
|
let size = shape[0] * shape[1] * shape[2];
|
|
let spacing = [1.0, 1.0, 1.0];
|
|
let origin = [0.0, 0.0, 0.0];
|
|
let affine = transform::from_spacing_origin(spacing, origin);
|
|
|
|
Self {
|
|
data: vec![0.0; size],
|
|
shape,
|
|
spacing,
|
|
origin,
|
|
affine,
|
|
}
|
|
}
|
|
|
|
/// Create a volume filled with a constant value.
|
|
pub fn filled(shape: [usize; 3], value: f64) -> Self {
|
|
let size = shape[0] * shape[1] * shape[2];
|
|
let spacing = [1.0, 1.0, 1.0];
|
|
let origin = [0.0, 0.0, 0.0];
|
|
let affine = transform::from_spacing_origin(spacing, origin);
|
|
|
|
Self {
|
|
data: vec![value; size],
|
|
shape,
|
|
spacing,
|
|
origin,
|
|
affine,
|
|
}
|
|
}
|
|
|
|
/// Get the volume dimensions [x, y, z].
|
|
pub fn shape(&self) -> [usize; 3] {
|
|
self.shape
|
|
}
|
|
|
|
/// Get the voxel spacing in mm [dx, dy, dz].
|
|
pub fn spacing(&self) -> [f64; 3] {
|
|
self.spacing
|
|
}
|
|
|
|
/// Get the origin in world coordinates [ox, oy, oz].
|
|
pub fn origin(&self) -> [f64; 3] {
|
|
self.origin
|
|
}
|
|
|
|
/// Get the affine transformation matrix.
|
|
pub fn affine(&self) -> Affine4 {
|
|
self.affine
|
|
}
|
|
|
|
/// Get a reference to the raw voxel data.
|
|
pub fn data(&self) -> &[f64] {
|
|
&self.data
|
|
}
|
|
|
|
/// Get a mutable reference to the raw voxel data.
|
|
pub fn data_mut(&mut self) -> &mut [f64] {
|
|
&mut self.data
|
|
}
|
|
|
|
/// Get the total number of voxels.
|
|
pub fn num_voxels(&self) -> usize {
|
|
self.shape[0] * self.shape[1] * self.shape[2]
|
|
}
|
|
|
|
/// Convert flat index to 3D coordinates (x, y, z).
|
|
pub fn index_to_coords(&self, index: usize) -> [usize; 3] {
|
|
let x = index % self.shape[0];
|
|
let y = (index / self.shape[0]) % self.shape[1];
|
|
let z = index / (self.shape[0] * self.shape[1]);
|
|
[x, y, z]
|
|
}
|
|
|
|
/// Convert 3D coordinates to flat index.
|
|
pub fn coords_to_index(&self, x: usize, y: usize, z: usize) -> usize {
|
|
z * (self.shape[0] * self.shape[1]) + y * self.shape[0] + x
|
|
}
|
|
|
|
/// Get voxel value at coordinates (x, y, z).
|
|
///
|
|
/// Returns None if coordinates are out of bounds.
|
|
pub fn get(&self, x: usize, y: usize, z: usize) -> Option<f64> {
|
|
if x < self.shape[0] && y < self.shape[1] && z < self.shape[2] {
|
|
Some(self.data[self.coords_to_index(x, y, z)])
|
|
} else {
|
|
None
|
|
}
|
|
}
|
|
|
|
/// Set voxel value at coordinates (x, y, z).
|
|
///
|
|
/// Returns false if coordinates are out of bounds.
|
|
pub fn set(&mut self, x: usize, y: usize, z: usize, value: f64) -> bool {
|
|
if x < self.shape[0] && y < self.shape[1] && z < self.shape[2] {
|
|
let idx = self.coords_to_index(x, y, z);
|
|
self.data[idx] = value;
|
|
true
|
|
} else {
|
|
false
|
|
}
|
|
}
|
|
|
|
/// Get voxel value at coordinates without bounds checking.
|
|
///
|
|
/// # Safety
|
|
/// Caller must ensure coordinates are valid.
|
|
#[allow(unsafe_op_in_unsafe_fn)]
|
|
pub unsafe fn get_unchecked(&self, x: usize, y: usize, z: usize) -> f64 {
|
|
let idx = self.coords_to_index(x, y, z);
|
|
// SAFETY: caller guarantees coordinates are valid
|
|
unsafe { *self.data.get_unchecked(idx) }
|
|
}
|
|
|
|
/// Set voxel value at coordinates without bounds checking.
|
|
///
|
|
/// # Safety
|
|
/// Caller must ensure coordinates are valid.
|
|
#[allow(unsafe_op_in_unsafe_fn)]
|
|
pub unsafe fn set_unchecked(&mut self, x: usize, y: usize, z: usize, value: f64) {
|
|
let idx = self.coords_to_index(x, y, z);
|
|
// SAFETY: caller guarantees coordinates are valid
|
|
unsafe { *self.data.get_unchecked_mut(idx) = value }
|
|
}
|
|
|
|
/// Convert voxel coordinates to world coordinates.
|
|
pub fn voxel_to_world(&self, voxel: [f64; 3]) -> [f64; 3] {
|
|
transform::voxel_to_world(&self.affine, voxel)
|
|
}
|
|
|
|
/// Convert world coordinates to voxel coordinates.
|
|
pub fn world_to_voxel(&self, world: [f64; 3]) -> Option<[f64; 3]> {
|
|
transform::world_to_voxel(&self.affine, world)
|
|
}
|
|
|
|
/// Get the bounding box in world coordinates.
|
|
///
|
|
/// Returns (min_corner, max_corner).
|
|
pub fn world_bounds(&self) -> ([f64; 3], [f64; 3]) {
|
|
let corners = [
|
|
[0.0, 0.0, 0.0],
|
|
[self.shape[0] as f64, 0.0, 0.0],
|
|
[0.0, self.shape[1] as f64, 0.0],
|
|
[0.0, 0.0, self.shape[2] as f64],
|
|
[self.shape[0] as f64, self.shape[1] as f64, 0.0],
|
|
[self.shape[0] as f64, 0.0, self.shape[2] as f64],
|
|
[0.0, self.shape[1] as f64, self.shape[2] as f64],
|
|
[
|
|
self.shape[0] as f64,
|
|
self.shape[1] as f64,
|
|
self.shape[2] as f64,
|
|
],
|
|
];
|
|
|
|
let world_corners: Vec<_> = corners.iter().map(|&c| self.voxel_to_world(c)).collect();
|
|
|
|
let min = [
|
|
world_corners
|
|
.iter()
|
|
.map(|c| c[0])
|
|
.fold(f64::INFINITY, f64::min),
|
|
world_corners
|
|
.iter()
|
|
.map(|c| c[1])
|
|
.fold(f64::INFINITY, f64::min),
|
|
world_corners
|
|
.iter()
|
|
.map(|c| c[2])
|
|
.fold(f64::INFINITY, f64::min),
|
|
];
|
|
|
|
let max = [
|
|
world_corners
|
|
.iter()
|
|
.map(|c| c[0])
|
|
.fold(f64::NEG_INFINITY, f64::max),
|
|
world_corners
|
|
.iter()
|
|
.map(|c| c[1])
|
|
.fold(f64::NEG_INFINITY, f64::max),
|
|
world_corners
|
|
.iter()
|
|
.map(|c| c[2])
|
|
.fold(f64::NEG_INFINITY, f64::max),
|
|
];
|
|
|
|
(min, max)
|
|
}
|
|
|
|
/// Get minimum voxel value.
|
|
pub fn min(&self) -> f64 {
|
|
self.data
|
|
.par_iter()
|
|
.copied()
|
|
.reduce(|| f64::INFINITY, f64::min)
|
|
}
|
|
|
|
/// Get maximum voxel value.
|
|
pub fn max(&self) -> f64 {
|
|
self.data
|
|
.par_iter()
|
|
.copied()
|
|
.reduce(|| f64::NEG_INFINITY, f64::max)
|
|
}
|
|
|
|
/// Get mean voxel value.
|
|
pub fn mean(&self) -> f64 {
|
|
let sum: f64 = self.data.par_iter().sum();
|
|
sum / self.data.len() as f64
|
|
}
|
|
|
|
/// Get standard deviation of voxel values.
|
|
pub fn std(&self) -> f64 {
|
|
let mean = self.mean();
|
|
let variance: f64 = self
|
|
.data
|
|
.par_iter()
|
|
.map(|&v| (v - mean).powi(2))
|
|
.sum::<f64>()
|
|
/ self.data.len() as f64;
|
|
variance.sqrt()
|
|
}
|
|
|
|
/// Get unique label values (for segmentation volumes).
|
|
pub fn unique_labels(&self) -> Vec<i64> {
|
|
use std::collections::BTreeSet;
|
|
let labels: BTreeSet<i64> = self.data.par_iter().map(|&v| v.round() as i64).collect();
|
|
labels.into_iter().collect()
|
|
}
|
|
|
|
/// Apply a function to each voxel value (in-place).
|
|
pub fn apply<F>(&mut self, f: F)
|
|
where
|
|
F: Fn(f64) -> f64 + Sync + Send,
|
|
{
|
|
self.data.par_iter_mut().for_each(|v| *v = f(*v));
|
|
}
|
|
|
|
/// Apply a function to each voxel, creating a new volume.
|
|
pub fn map<F>(&self, f: F) -> Self
|
|
where
|
|
F: Fn(f64) -> f64 + Sync + Send,
|
|
{
|
|
let new_data: Vec<f64> = self.data.par_iter().map(|&v| f(v)).collect();
|
|
Self {
|
|
data: new_data,
|
|
shape: self.shape,
|
|
spacing: self.spacing,
|
|
origin: self.origin,
|
|
affine: self.affine,
|
|
}
|
|
}
|
|
|
|
/// Threshold the volume: values >= threshold become 1.0, others become 0.0.
|
|
pub fn threshold(&self, threshold: f64) -> Self {
|
|
self.map(|v| if v >= threshold { 1.0 } else { 0.0 })
|
|
}
|
|
|
|
/// Create a binary mask where values equal to the label are 1.0.
|
|
pub fn mask_label(&self, label: i64) -> Self {
|
|
self.map(|v| if v.round() as i64 == label { 1.0 } else { 0.0 })
|
|
}
|
|
|
|
/// Extract a sub-volume (crop).
|
|
///
|
|
/// # Arguments
|
|
/// * `start` - Start coordinates [x, y, z]
|
|
/// * `size` - Size of sub-volume [sx, sy, sz]
|
|
pub fn crop(&self, start: [usize; 3], size: [usize; 3]) -> Self {
|
|
let [sx, sy, sz] = size;
|
|
let mut new_data = vec![0.0; sx * sy * sz];
|
|
|
|
for z in 0..sz {
|
|
for y in 0..sy {
|
|
for x in 0..sx {
|
|
let src_idx = self.coords_to_index(start[0] + x, start[1] + y, start[2] + z);
|
|
let dst_idx = z * (sx * sy) + y * sx + x;
|
|
new_data[dst_idx] = self.data[src_idx];
|
|
}
|
|
}
|
|
}
|
|
|
|
// Adjust origin for the crop
|
|
let new_origin = self.voxel_to_world([start[0] as f64, start[1] as f64, start[2] as f64]);
|
|
|
|
Self {
|
|
data: new_data,
|
|
shape: size,
|
|
spacing: self.spacing,
|
|
origin: new_origin,
|
|
affine: self.affine, // Affine stays the same, origin adjusted
|
|
}
|
|
}
|
|
|
|
/// Set spacing and update the affine accordingly.
|
|
pub fn set_spacing(&mut self, spacing: [f64; 3]) {
|
|
self.spacing = spacing;
|
|
self.affine = transform::from_spacing_origin(spacing, self.origin);
|
|
}
|
|
|
|
/// Set origin and update the affine accordingly.
|
|
pub fn set_origin(&mut self, origin: [f64; 3]) {
|
|
self.origin = origin;
|
|
self.affine = transform::from_spacing_origin(self.spacing, origin);
|
|
}
|
|
|
|
/// Set the affine transformation and extract spacing/origin.
|
|
pub fn set_affine(&mut self, affine: Affine4) {
|
|
self.affine = affine;
|
|
self.spacing = transform::get_spacing(&affine);
|
|
self.origin = transform::get_origin(&affine);
|
|
}
|
|
|
|
/// Trilinear interpolation at a floating-point voxel coordinate.
|
|
///
|
|
/// Returns None if the coordinate is outside the volume.
|
|
pub fn interpolate(&self, x: f64, y: f64, z: f64) -> Option<f64> {
|
|
if x < 0.0
|
|
|| y < 0.0
|
|
|| z < 0.0
|
|
|| x >= (self.shape[0] - 1) as f64
|
|
|| y >= (self.shape[1] - 1) as f64
|
|
|| z >= (self.shape[2] - 1) as f64
|
|
{
|
|
return None;
|
|
}
|
|
|
|
let x0 = x.floor() as usize;
|
|
let y0 = y.floor() as usize;
|
|
let z0 = z.floor() as usize;
|
|
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;
|
|
|
|
// Get 8 corner values
|
|
let c000 = self.get(x0, y0, z0)?;
|
|
let c100 = self.get(x1, y0, z0)?;
|
|
let c010 = self.get(x0, y1, z0)?;
|
|
let c110 = self.get(x1, y1, z0)?;
|
|
let c001 = self.get(x0, y0, z1)?;
|
|
let c101 = self.get(x1, y0, z1)?;
|
|
let c011 = self.get(x0, y1, z1)?;
|
|
let c111 = self.get(x1, y1, z1)?;
|
|
|
|
// 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
|
|
Some(c0 * (1.0 - zd) + c1 * zd)
|
|
}
|
|
|
|
/// Sample the volume at a world coordinate using trilinear interpolation.
|
|
pub fn sample_world(&self, world: [f64; 3]) -> Option<f64> {
|
|
let voxel = self.world_to_voxel(world)?;
|
|
self.interpolate(voxel[0], voxel[1], voxel[2])
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
#[test]
|
|
fn test_zeros() {
|
|
let vol = Volume::zeros([10, 20, 30]);
|
|
assert_eq!(vol.shape(), [10, 20, 30]);
|
|
assert_eq!(vol.num_voxels(), 6000);
|
|
assert_eq!(vol.min(), 0.0);
|
|
assert_eq!(vol.max(), 0.0);
|
|
}
|
|
|
|
#[test]
|
|
fn test_get_set() {
|
|
let mut vol = Volume::zeros([10, 10, 10]);
|
|
vol.set(5, 5, 5, 42.0);
|
|
assert_eq!(vol.get(5, 5, 5), Some(42.0));
|
|
assert_eq!(vol.get(0, 0, 0), Some(0.0));
|
|
assert_eq!(vol.get(100, 100, 100), None);
|
|
}
|
|
|
|
#[test]
|
|
fn test_coords_conversion() {
|
|
let vol = Volume::zeros([10, 20, 30]);
|
|
|
|
// Test round-trip
|
|
for z in 0..5 {
|
|
for y in 0..5 {
|
|
for x in 0..5 {
|
|
let idx = vol.coords_to_index(x, y, z);
|
|
let [rx, ry, rz] = vol.index_to_coords(idx);
|
|
assert_eq!([x, y, z], [rx, ry, rz]);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn test_statistics() {
|
|
let vol = Volume::filled([10, 10, 10], 5.0);
|
|
assert_eq!(vol.mean(), 5.0);
|
|
assert_eq!(vol.std(), 0.0);
|
|
}
|
|
|
|
#[test]
|
|
fn test_interpolate() {
|
|
let mut vol = Volume::zeros([10, 10, 10]);
|
|
vol.set(0, 0, 0, 0.0);
|
|
vol.set(1, 0, 0, 1.0);
|
|
|
|
// Interpolate at midpoint
|
|
let val = vol.interpolate(0.5, 0.0, 0.0).unwrap();
|
|
assert!((val - 0.5).abs() < 1e-10);
|
|
}
|
|
}
|