Initial commit
This commit is contained in:
redclawsystems
@@ -0,0 +1,87 @@
|
||||
//! Similarity metrics for image registration.
|
||||
|
||||
mod mse;
|
||||
mod ncc;
|
||||
|
||||
pub use mse::MeanSquaredError;
|
||||
pub use ncc::NormalizedCrossCorrelation;
|
||||
|
||||
use rtx_medical_io::Volume;
|
||||
|
||||
/// Trait for image similarity metrics.
|
||||
pub trait SimilarityMetric: Send + Sync {
|
||||
/// Compute the similarity between two volumes.
|
||||
///
|
||||
/// Returns a value where:
|
||||
/// - Lower values indicate better alignment (for MSE-like metrics)
|
||||
/// - Higher values indicate better alignment (for correlation-like metrics)
|
||||
fn compute(&self, fixed: &Volume, moving: &Volume, mask: Option<&Volume>) -> f64;
|
||||
|
||||
/// Return true if higher values mean better alignment.
|
||||
fn is_similarity(&self) -> bool;
|
||||
|
||||
/// Convert to a cost (lower is better).
|
||||
fn to_cost(&self, value: f64) -> f64 {
|
||||
if self.is_similarity() { -value } else { value }
|
||||
}
|
||||
}
|
||||
|
||||
/// Compute the gradient of a volume at each voxel.
|
||||
pub fn compute_gradient(volume: &Volume) -> (Volume, Volume, Volume) {
|
||||
let shape = volume.shape();
|
||||
let spacing = volume.spacing();
|
||||
|
||||
let mut grad_x = Volume::zeros(shape);
|
||||
let mut grad_y = Volume::zeros(shape);
|
||||
let mut grad_z = Volume::zeros(shape);
|
||||
|
||||
for z in 1..shape[2] - 1 {
|
||||
for y in 1..shape[1] - 1 {
|
||||
for x in 1..shape[0] - 1 {
|
||||
// Central differences
|
||||
let v_xp = volume.get(x + 1, y, z).unwrap_or(0.0);
|
||||
let v_xm = volume.get(x - 1, y, z).unwrap_or(0.0);
|
||||
let v_yp = volume.get(x, y + 1, z).unwrap_or(0.0);
|
||||
let v_ym = volume.get(x, y - 1, z).unwrap_or(0.0);
|
||||
let v_zp = volume.get(x, y, z + 1).unwrap_or(0.0);
|
||||
let v_zm = volume.get(x, y, z - 1).unwrap_or(0.0);
|
||||
|
||||
let dx = (v_xp - v_xm) / (2.0 * spacing[0]);
|
||||
let dy = (v_yp - v_ym) / (2.0 * spacing[1]);
|
||||
let dz = (v_zp - v_zm) / (2.0 * spacing[2]);
|
||||
|
||||
grad_x.set(x, y, z, dx);
|
||||
grad_y.set(x, y, z, dy);
|
||||
grad_z.set(x, y, z, dz);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
(grad_x, grad_y, grad_z)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_gradient() {
|
||||
let mut vol = Volume::zeros([10, 10, 10]);
|
||||
|
||||
// Create a linear gradient in X direction
|
||||
for z in 0..10 {
|
||||
for y in 0..10 {
|
||||
for x in 0..10 {
|
||||
vol.set(x, y, z, x as f64);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let (grad_x, grad_y, grad_z) = compute_gradient(&vol);
|
||||
|
||||
// Interior points should have gradient ~1 in X, ~0 in Y and Z
|
||||
assert!((grad_x.get(5, 5, 5).unwrap() - 1.0).abs() < 1e-10);
|
||||
assert!(grad_y.get(5, 5, 5).unwrap().abs() < 1e-10);
|
||||
assert!(grad_z.get(5, 5, 5).unwrap().abs() < 1e-10);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,146 @@
|
||||
//! Mean Squared Error similarity metric.
|
||||
|
||||
use super::SimilarityMetric;
|
||||
use rayon::prelude::*;
|
||||
use rtx_medical_io::Volume;
|
||||
|
||||
/// Mean Squared Error (MSE) metric.
|
||||
///
|
||||
/// Computes: MSE = (1/N) * Σ(fixed[i] - moving[i])²
|
||||
///
|
||||
/// Lower values indicate better alignment.
|
||||
#[derive(Debug, Clone, Default)]
|
||||
pub struct MeanSquaredError;
|
||||
|
||||
impl MeanSquaredError {
|
||||
/// Create a new MSE metric.
|
||||
pub fn new() -> Self {
|
||||
Self
|
||||
}
|
||||
}
|
||||
|
||||
impl SimilarityMetric for MeanSquaredError {
|
||||
fn compute(&self, fixed: &Volume, moving: &Volume, mask: Option<&Volume>) -> f64 {
|
||||
let shape = fixed.shape();
|
||||
|
||||
if shape != moving.shape() {
|
||||
return f64::MAX;
|
||||
}
|
||||
|
||||
// Collect all voxel pairs
|
||||
let total_voxels = shape[0] * shape[1] * shape[2];
|
||||
let results: Vec<(f64, usize)> = (0..total_voxels)
|
||||
.into_par_iter()
|
||||
.map(|idx| {
|
||||
let x = idx % shape[0];
|
||||
let y = (idx / shape[0]) % shape[1];
|
||||
let z = idx / (shape[0] * shape[1]);
|
||||
|
||||
// Check mask
|
||||
if let Some(m) = mask {
|
||||
if m.get(x, y, z).unwrap_or(0.0) <= 0.0 {
|
||||
return (0.0, 0);
|
||||
}
|
||||
}
|
||||
|
||||
let f = fixed.get(x, y, z).unwrap_or(0.0);
|
||||
let m = moving.get(x, y, z).unwrap_or(0.0);
|
||||
let diff = f - m;
|
||||
(diff * diff, 1)
|
||||
})
|
||||
.collect();
|
||||
|
||||
let (sum_sq, count): (f64, usize) = results
|
||||
.iter()
|
||||
.fold((0.0, 0), |acc, &(sq, c)| (acc.0 + sq, acc.1 + c));
|
||||
|
||||
if count == 0 {
|
||||
return f64::MAX;
|
||||
}
|
||||
|
||||
sum_sq / count as f64
|
||||
}
|
||||
|
||||
fn is_similarity(&self) -> bool {
|
||||
false // Lower is better
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_identical_images() {
|
||||
let mut vol = Volume::zeros([10, 10, 10]);
|
||||
for z in 0..10 {
|
||||
for y in 0..10 {
|
||||
for x in 0..10 {
|
||||
vol.set(x, y, z, (x + y + z) as f64);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let mse = MeanSquaredError::new();
|
||||
let value = mse.compute(&vol, &vol, None);
|
||||
assert!(value < 1e-10);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_constant_difference() {
|
||||
let mut vol1 = Volume::zeros([10, 10, 10]);
|
||||
let mut vol2 = Volume::zeros([10, 10, 10]);
|
||||
|
||||
for z in 0..10 {
|
||||
for y in 0..10 {
|
||||
for x in 0..10 {
|
||||
vol1.set(x, y, z, 5.0);
|
||||
vol2.set(x, y, z, 7.0); // Constant difference of 2
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let mse = MeanSquaredError::new();
|
||||
let value = mse.compute(&vol1, &vol2, None);
|
||||
assert!((value - 4.0).abs() < 1e-10); // MSE = (2)² = 4
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_with_mask() {
|
||||
let mut vol1 = Volume::zeros([10, 10, 10]);
|
||||
let mut vol2 = Volume::zeros([10, 10, 10]);
|
||||
let mut mask = Volume::zeros([10, 10, 10]);
|
||||
|
||||
// Only set mask for a small region
|
||||
for z in 0..5 {
|
||||
for y in 0..5 {
|
||||
for x in 0..5 {
|
||||
vol1.set(x, y, z, 5.0);
|
||||
vol2.set(x, y, z, 7.0);
|
||||
mask.set(x, y, z, 1.0);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Rest of volumes have different values but should be ignored
|
||||
for z in 5..10 {
|
||||
for y in 5..10 {
|
||||
for x in 5..10 {
|
||||
vol1.set(x, y, z, 100.0);
|
||||
vol2.set(x, y, z, 0.0);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let mse = MeanSquaredError::new();
|
||||
let value = mse.compute(&vol1, &vol2, Some(&mask));
|
||||
assert!((value - 4.0).abs() < 1e-10); // MSE should only consider masked region
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_to_cost() {
|
||||
let mse = MeanSquaredError::new();
|
||||
assert!(!mse.is_similarity());
|
||||
assert_eq!(mse.to_cost(5.0), 5.0); // Already a cost
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,191 @@
|
||||
//! Normalized Cross-Correlation similarity metric.
|
||||
|
||||
use super::SimilarityMetric;
|
||||
use rayon::prelude::*;
|
||||
use rtx_medical_io::Volume;
|
||||
|
||||
/// Normalized Cross-Correlation (NCC) metric.
|
||||
///
|
||||
/// Computes: NCC = Σ((f - μf)(m - μm)) / (σf * σm * N)
|
||||
///
|
||||
/// Returns values in [-1, 1], where 1 indicates perfect correlation.
|
||||
#[derive(Debug, Clone, Default)]
|
||||
pub struct NormalizedCrossCorrelation;
|
||||
|
||||
impl NormalizedCrossCorrelation {
|
||||
/// Create a new NCC metric.
|
||||
pub fn new() -> Self {
|
||||
Self
|
||||
}
|
||||
}
|
||||
|
||||
impl SimilarityMetric for NormalizedCrossCorrelation {
|
||||
fn compute(&self, fixed: &Volume, moving: &Volume, mask: Option<&Volume>) -> f64 {
|
||||
let shape = fixed.shape();
|
||||
|
||||
if shape != moving.shape() {
|
||||
return -1.0;
|
||||
}
|
||||
|
||||
let total_voxels = shape[0] * shape[1] * shape[2];
|
||||
|
||||
// First pass: compute means
|
||||
let stats: Vec<(f64, f64, usize)> = (0..total_voxels)
|
||||
.into_par_iter()
|
||||
.map(|idx| {
|
||||
let x = idx % shape[0];
|
||||
let y = (idx / shape[0]) % shape[1];
|
||||
let z = idx / (shape[0] * shape[1]);
|
||||
|
||||
if let Some(m) = mask {
|
||||
if m.get(x, y, z).unwrap_or(0.0) <= 0.0 {
|
||||
return (0.0, 0.0, 0);
|
||||
}
|
||||
}
|
||||
|
||||
let f = fixed.get(x, y, z).unwrap_or(0.0);
|
||||
let m = moving.get(x, y, z).unwrap_or(0.0);
|
||||
(f, m, 1)
|
||||
})
|
||||
.collect();
|
||||
|
||||
let (sum_f, sum_m, count): (f64, f64, usize) =
|
||||
stats.iter().fold((0.0, 0.0, 0), |acc, &(f, m, c)| {
|
||||
(acc.0 + f, acc.1 + m, acc.2 + c)
|
||||
});
|
||||
|
||||
if count == 0 {
|
||||
return -1.0;
|
||||
}
|
||||
|
||||
let mean_f = sum_f / count as f64;
|
||||
let mean_m = sum_m / count as f64;
|
||||
|
||||
// Second pass: compute correlation and variances
|
||||
let corr_stats: Vec<(f64, f64, f64)> = (0..total_voxels)
|
||||
.into_par_iter()
|
||||
.map(|idx| {
|
||||
let x = idx % shape[0];
|
||||
let y = (idx / shape[0]) % shape[1];
|
||||
let z = idx / (shape[0] * shape[1]);
|
||||
|
||||
if let Some(m) = mask {
|
||||
if m.get(x, y, z).unwrap_or(0.0) <= 0.0 {
|
||||
return (0.0, 0.0, 0.0);
|
||||
}
|
||||
}
|
||||
|
||||
let f = fixed.get(x, y, z).unwrap_or(0.0) - mean_f;
|
||||
let m = moving.get(x, y, z).unwrap_or(0.0) - mean_m;
|
||||
(f * m, f * f, m * m)
|
||||
})
|
||||
.collect();
|
||||
|
||||
let (sum_fm, sum_ff, sum_mm): (f64, f64, f64) = corr_stats
|
||||
.iter()
|
||||
.fold((0.0, 0.0, 0.0), |acc, &(fm, ff, mm)| {
|
||||
(acc.0 + fm, acc.1 + ff, acc.2 + mm)
|
||||
});
|
||||
|
||||
let std_f = sum_ff.sqrt();
|
||||
let std_m = sum_mm.sqrt();
|
||||
|
||||
if std_f < 1e-10 || std_m < 1e-10 {
|
||||
return 0.0;
|
||||
}
|
||||
|
||||
sum_fm / (std_f * std_m)
|
||||
}
|
||||
|
||||
fn is_similarity(&self) -> bool {
|
||||
true // Higher is better
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_identical_images() {
|
||||
let mut vol = Volume::zeros([10, 10, 10]);
|
||||
for z in 0..10 {
|
||||
for y in 0..10 {
|
||||
for x in 0..10 {
|
||||
vol.set(x, y, z, (x * y + z) as f64);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let ncc = NormalizedCrossCorrelation::new();
|
||||
let value = ncc.compute(&vol, &vol, None);
|
||||
assert!((value - 1.0).abs() < 1e-10);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_scaled_images() {
|
||||
let mut vol1 = Volume::zeros([10, 10, 10]);
|
||||
let mut vol2 = Volume::zeros([10, 10, 10]);
|
||||
|
||||
for z in 0..10 {
|
||||
for y in 0..10 {
|
||||
for x in 0..10 {
|
||||
let v = (x + y + z) as f64 + 1.0; // Avoid zero variance
|
||||
vol1.set(x, y, z, v);
|
||||
vol2.set(x, y, z, v * 2.0); // Scaled version
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let ncc = NormalizedCrossCorrelation::new();
|
||||
let value = ncc.compute(&vol1, &vol2, None);
|
||||
assert!((value - 1.0).abs() < 1e-10); // NCC is invariant to scaling
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_shifted_images() {
|
||||
let mut vol1 = Volume::zeros([10, 10, 10]);
|
||||
let mut vol2 = Volume::zeros([10, 10, 10]);
|
||||
|
||||
for z in 0..10 {
|
||||
for y in 0..10 {
|
||||
for x in 0..10 {
|
||||
let v = (x + y + z) as f64 + 1.0;
|
||||
vol1.set(x, y, z, v);
|
||||
vol2.set(x, y, z, v + 100.0); // Shifted version
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let ncc = NormalizedCrossCorrelation::new();
|
||||
let value = ncc.compute(&vol1, &vol2, None);
|
||||
assert!((value - 1.0).abs() < 1e-10); // NCC is invariant to shifting
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_negative_correlation() {
|
||||
let mut vol1 = Volume::zeros([10, 10, 10]);
|
||||
let mut vol2 = Volume::zeros([10, 10, 10]);
|
||||
|
||||
for z in 0..10 {
|
||||
for y in 0..10 {
|
||||
for x in 0..10 {
|
||||
let v = (x + y + z) as f64 + 1.0;
|
||||
vol1.set(x, y, z, v);
|
||||
vol2.set(x, y, z, -v); // Negated version
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let ncc = NormalizedCrossCorrelation::new();
|
||||
let value = ncc.compute(&vol1, &vol2, None);
|
||||
assert!((value + 1.0).abs() < 1e-10); // Should be -1
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_to_cost() {
|
||||
let ncc = NormalizedCrossCorrelation::new();
|
||||
assert!(ncc.is_similarity());
|
||||
assert_eq!(ncc.to_cost(0.8), -0.8); // Negate for cost
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user