147 lines
3.9 KiB
Rust
147 lines
3.9 KiB
Rust
//! 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
|
|
}
|
|
}
|