wim/image-registration
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

- cleanup

Browse Source
This commit is contained in:
w.pomp
2026-07-31 17:49:11 +02:00
parent f75f12db16
commit cc12a1a9b9
5 changed files with 253 additions and 175 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/filter.rs
-1
View File
@@ -142,7 +142,6 @@ mod tests {
Ok(())
}
#[test]
fn fft_test() -> Result<(), Box<dyn std::error::Error>> {
let x = Array1::linspace(0.0, 1.0, 11).to_vec();
src/metric.rs
+5 -23
View File
@@ -152,7 +152,7 @@ fn parzen(
enum Sampling {
Fixed(Vec<Vec<f64>>),
Random((usize, Vec<f64>, RefCell<Vec<Vec<f64>>>)),
Random((usize, Vec<f64>)),
}
impl Sampling {
@@ -161,25 +161,19 @@ impl Sampling {
}
fn random(n_samples: usize, shape: Vec<f64>) -> Self {
Self::Random((n_samples, shape, RefCell::new(Vec::new())))
Self::Random((n_samples, shape))
}
fn index(&self) -> Vec<Vec<f64>> {
match self {
Self::Fixed(index) => index.clone(),
Self::Random((n_samples, shape, cached)) => {
let mut cache = cached.borrow_mut();
if cache.is_empty() {
*cache = rand::rng()
Self::Random((n_samples, shape)) => rand::rng()
.random_iter::<f64>()
.take(n_samples * shape.len())
.chunks(shape.len())
.into_iter()
.map(|c| c.zip_eq(shape.iter()).map(|(i, s)| i * s - 0.5).collect())
.collect();
}
cache.clone()
}
.collect(),
}
}
}
@@ -437,23 +431,11 @@ where
debug_assert!(metric.is_finite());
debug_assert!(dmetric.iter().all(|i| i.is_finite()));
// let alpha = 1.0;
// let dalpha = vec![0.0; dalpha.len()];
self.metric.replace(IntMut {
mu: mu.to_vec(),
metric,
derivative: dmetric.clone(),
});
// self.metric.replace(IntMut {
// mu: mu.to_vec(),
// metric: metric / alpha,
// derivative: dmetric
// .iter()
// .zip_eq(dalpha)
// .map(|(dm, da)| *dm / alpha + da * metric / (alpha * alpha))
// .collect(),
// });
}
Ok(())
}
@@ -778,7 +760,7 @@ mod tests {
let fixed = BSpline::<0, _>::new(a.view());
let moving = BSpline::<3, _>::new(b.view());
let mus = vec![0.0, 0.001];
let mus = [0.0, 0.001];
let mut npz = NpzWriter::new(File::create(
std::env::home_dir().unwrap().join("tmp/metric.npz"),
)?);
src/optimize.rs
+14 -14
View File
@@ -358,6 +358,20 @@ mod tests {
}
}
struct QuadraticWithMinimum;
impl ObjectiveFunction<f64> for QuadraticWithMinimum {
fn evaluate(&self, point: &[f64]) -> f64 {
let x = point[0];
(x - 2.0).powi(2)
}
fn gradient(&self, point: &[f64]) -> Option<Vec<f64>> {
let x = point[0];
Some(vec![2.0 * (x - 2.0)])
}
}
#[test]
fn test_lbfgs_quadratic() {
let f = Quadratic;
@@ -379,20 +393,6 @@ mod tests {
#[test]
fn test_lbfgs_quadratic_with_minimum() {
struct QuadraticWithMinimum;
impl ObjectiveFunction<f64> for QuadraticWithMinimum {
fn evaluate(&self, point: &[f64]) -> f64 {
let x = point[0];
(x - 2.0).powi(2)
}
fn gradient(&self, point: &[f64]) -> Option<Vec<f64>> {
let x = point[0];
Some(vec![2.0 * (x - 2.0)])
}
}
let f = QuadraticWithMinimum;
let initial_point = vec![0.0];
let config = OptimizationConfig {
src/register.rs
+55 -108
View File
@@ -9,8 +9,7 @@ use num::cast::AsPrimitive;
use std::marker::PhantomData;
/// Optimizer type for registration
#[derive(Clone, Debug, PartialEq)]
#[derive(Default)]
#[derive(Clone, Debug, PartialEq, Default)]
pub enum Optimizer {
/// L-BFGS optimizer (fast, requires consistent gradients)
#[default]
@@ -19,7 +18,6 @@ pub enum Optimizer {
ASGD,
}
#[derive(Clone, Debug)]
pub struct RegistrationStep {
pub sigma: Sigma,
@@ -29,7 +27,6 @@ pub struct RegistrationStep {
pub edge: f64,
pub max_iterations: usize,
pub learning_rate: f64,
pub downsample: usize,
/// Optimizer to use for this level
pub optimizer: Optimizer,
}
@@ -52,112 +49,64 @@ impl RegistrationStep {
edge,
max_iterations,
learning_rate,
downsample: 1,
optimizer: Optimizer::default(),
}
}
/// Default registration steps matching elastix `FixedSmoothingImagePyramid`.
///
/// Uses a hybrid multi-fidelity approach:
/// - Coarse levels (sigma=4,2): ASGD optimizer with all pixels for speed
/// - Fine levels (sigma=1,0.5): L-BFGS optimizer with all pixels for accuracy
///
/// Note: Samples are cached per level (not randomized per iteration like SimpleElastix),
/// but ASGD's sigmoid momentum handles the stochastic gradients effectively.
pub fn default_steps(ndim: usize, n: usize) -> Vec<Self> {
if ndim == 1 {
return vec![Self {
sigma: Sigma::Absolute(vec![1.0; 1]),
samples: SamplingArg::Random(n),
vec![
Self {
sigma: Sigma::Absolute(vec![8.0; ndim]),
samples: SamplingArg::Random((n / 256).max(1024)),
n_bins: 16,
tolerance: 1e-1,
edge: 0.05,
max_iterations: 128,
learning_rate: 1.0,
optimizer: Optimizer::ASGD,
},
Self {
sigma: Sigma::Absolute(vec![4.0; ndim]),
samples: SamplingArg::Random((n / 64).max(4096)),
n_bins: 24,
tolerance: 1e-2,
edge: 0.05,
max_iterations: 256,
learning_rate: 1.0,
optimizer: Optimizer::ASGD,
},
Self {
sigma: Sigma::Absolute(vec![2.0; ndim]),
samples: SamplingArg::Random((n / 16).max(4096)),
n_bins: 32,
tolerance: 1e-4,
edge: 0.05,
max_iterations: 512,
learning_rate: 1.0,
optimizer: Optimizer::ASGD,
},
Self {
sigma: Sigma::Absolute(vec![1.0; ndim]),
samples: SamplingArg::Random((n / 4).max(8192)),
n_bins: 48,
tolerance: 1e-6,
edge: 0.05,
max_iterations: 1024,
learning_rate: 1.0,
optimizer: Optimizer::ASGD,
},
Self {
sigma: Sigma::None,
samples: SamplingArg::Fixed(n.max(16384)),
n_bins: 64,
tolerance: 1e-8,
edge: 0.05,
max_iterations: 2048,
learning_rate: 1.0,
downsample: 1,
optimizer: Optimizer::LBFGS,
}];
},
]
}
// Pyramid with downsampling: schedule [8,4,2,1], sigma [4,2,1,0.5]
let sigma_schedule: Vec<f64> = vec![4.0, 2.0, 1.0, 0.5];
let downsample_schedule: Vec<usize> = vec![8, 4, 2, 1];
let n_levels = sigma_schedule.len();
sigma_schedule
.iter()
.zip(downsample_schedule.iter())
.enumerate()
.map(|(level, (&s, &d))| {
let n_pixels = (n / (d * d)).max(4);
let is_finest = level == n_levels - 1;
let is_coarse = level < 2;
// Hybrid: ASGD on coarse levels for speed, L-BFGS on fine levels for accuracy
// Use all pixels on all levels for consistent gradient estimates
let (optimizer, max_iter, tol) = if is_coarse {
(Optimizer::ASGD, 512, 1e-4)
} else if is_finest {
(Optimizer::LBFGS, 2048, 1e-8)
} else {
(Optimizer::LBFGS, 512, 1e-6)
};
Self {
sigma: Sigma::Absolute(vec![s; ndim]),
samples: SamplingArg::Random(n_pixels),
n_bins: 32,
tolerance: tol,
edge: 0.05,
max_iterations: max_iter,
learning_rate: 1.0,
downsample: d,
optimizer,
}
})
.collect()
}
}
/// Downsample an n-dimensional array by integer factor along all axes.
/// Takes every `factor`-th element along each axis.
/// Only supports 1D and 2D arrays (the only types used in this crate).
pub fn downsample_nd<D>(array: ndarray::ArrayView<f64, D>, factor: usize) -> ndarray::Array<f64, D>
where
D: ndarray::Dimension,
{
if factor <= 1 {
return array.to_owned();
}
let ndim = array.ndim();
let src = array.to_owned();
let src_shape = src.shape().to_vec();
let src_slice = src.as_slice().unwrap();
let new_shape: Vec<usize> = src_shape.iter().map(|&s| (s - 1) / factor + 1).collect();
// Compute strides for row-major layout
let mut strides = vec![1usize; ndim];
for i in (0..ndim - 1).rev() {
strides[i] = strides[i + 1] * src_shape[i + 1];
}
let total: usize = new_shape.iter().product();
let mut data = Vec::with_capacity(total);
for out_flat in 0..total {
let mut out_coords = vec![0usize; ndim];
let mut tmp = out_flat;
for i in 0..ndim {
out_coords[i] = tmp % new_shape[i];
tmp /= new_shape[i];
}
let mut src_flat = 0;
for i in 0..ndim {
src_flat += out_coords[i] * factor * strides[i];
}
data.push(src_slice[src_flat]);
}
ndarray::Array::from_shape_vec(new_shape, data)
.unwrap()
.into_dimensionality::<D>()
.unwrap()
}
#[derive(Clone, Debug)]
@@ -283,15 +232,14 @@ impl<D: Dimension> Registration<D> {
fn golden_standard_scales(fixed_shape: &[usize], ndim: usize) -> Vec<f64> {
let n_params = ndim * ndim + ndim;
let mut scales = Vec::with_capacity(n_params);
for _i in 0..ndim {
for j in 0..ndim {
let n = fixed_shape[j] as f64;
scales.push(((n * n - 1.0) / 12.0).sqrt());
}
}
for _ in 0..ndim {
scales.push(1.0);
scales.extend(
fixed_shape
.iter()
.map(|&i| ((i * i - 1) as f64 / 12.0).sqrt()),
);
}
scales.extend(vec![1.0; ndim]);
scales
}
@@ -325,7 +273,6 @@ impl<D: Dimension> Registration<D> {
edge,
max_iterations,
learning_rate,
downsample: _,
optimizer,
} in steps.into_iter()
{
@@ -356,7 +303,8 @@ impl<D: Dimension> Registration<D> {
);
if let Some(result) = optimization_result
&& result.optimal_point.iter().all(|i| i.is_finite()) {
&& result.optimal_point.iter().all(|i| i.is_finite())
{
p = self.fixed_mu.combine(&result.optimal_point);
}
}
@@ -431,7 +379,6 @@ impl<D: Dimension> Registration<D> {
edge,
max_iterations,
learning_rate,
downsample: _,
optimizer,
} in steps.into_iter()
{
src/transform.rs
+171 -21
View File
@@ -579,7 +579,7 @@ mod tests {
use crate::julia_image;
use crate::transform::Transform;
use itertools::Itertools;
use ndarray::{Array2, Ix2, s};
use ndarray::{Array2, Ix2, array, s};
use num::traits::FloatConst;
use std::fs::File;
use std::path::Path;
@@ -789,11 +789,11 @@ mod tests {
// Test several points
let test_points: Vec<(&str, Vec<f64>)> = vec![
("identity".into(), vec![1.0, 0.0]),
("truth".into(), vec![0.85, 4.0]),
("neg_trans".into(), vec![1.0, -4.0]),
("scale_0.9".into(), vec![0.9, 0.0]),
("scale_1.1".into(), vec![1.1, 0.0]),
("identity", vec![1.0, 0.0]),
("truth", vec![0.85, 4.0]),
("neg_trans", vec![1.0, -4.0]),
("scale_0.9", vec![0.9, 0.0]),
("scale_1.1", vec![1.1, 0.0]),
];
println!("=== SMOOTHED ALL 100 integer points bins=32 ===");
@@ -831,6 +831,96 @@ mod tests {
#[test]
fn register1() -> Result<(), Box<dyn std::error::Error>> {
let fixed = array![
0.00133107, 0.00138621, 0.00141719, 0.00137673, 0.00133192, 0.00131586, 0.00129737,
0.00131649, 0.00132131, 0.0013368, 0.0013557, 0.00135557, 0.00135903, 0.00136225,
0.00136999, 0.0013847, 0.00137785, 0.00138041, 0.00140829, 0.00142085, 0.0014344,
0.00149646, 0.00156956, 0.00172352, 0.0019473, 0.00226109, 0.00269055, 0.00312467,
0.00343227, 0.00352425, 0.00335073, 0.00305988, 0.00276482, 0.0024797, 0.00223358,
0.00204837, 0.00191672, 0.00183362, 0.00179435, 0.00177568, 0.00177995, 0.00177444,
0.00178801, 0.00180613, 0.00181906, 0.00183161, 0.00185975, 0.00185648, 0.00187759,
0.00188436, 0.00185824, 0.00185461, 0.00188928, 0.00190146, 0.00185707, 0.00185682,
0.00185372, 0.00184588, 0.00183992, 0.0018354, 0.0018307, 0.00181903, 0.00182027,
0.00181646, 0.00180596, 0.00180369, 0.00180493, 0.0018042, 0.00180887, 0.00180575,
0.00180144, 0.00180303, 0.00180743, 0.00180482, 0.0017817, 0.00175849, 0.00174724,
0.00174177, 0.00174569, 0.00176198, 0.00178241, 0.00180581, 0.00184414, 0.00188772,
0.00192021, 0.00195737, 0.00201045, 0.00207885, 0.00216532, 0.00227495, 0.00240301,
0.00254983, 0.00273925, 0.00298603, 0.00329465, 0.00367626, 0.00415677, 0.0047952,
0.00565066, 0.006755, 0.00808893, 0.00958421, 0.0111741, 0.01279892, 0.01434347,
0.0156247, 0.01649142, 0.01690067, 0.01689863, 0.0165402, 0.0158715, 0.01497264,
0.0139421, 0.01286888, 0.01182453, 0.01084795, 0.00994079, 0.00908256, 0.00826134,
0.00748309, 0.00676281, 0.00611948, 0.00556657, 0.00510174, 0.00470918, 0.0043702,
0.00407155, 0.00381196, 0.00359745, 0.00342831, 0.00329565, 0.00318778, 0.00309373,
0.00300738, 0.00292954, 0.00286483, 0.00281634, 0.00278162, 0.00275487, 0.00273168,
0.0027102, 0.0026901, 0.00267163, 0.00265643, 0.00264636, 0.00264188, 0.00264216,
0.00264419, 0.00264316, 0.00263564, 0.00262231, 0.00260767, 0.00259683, 0.00259233,
0.00259301, 0.00259521, 0.0025956, 0.00259316, 0.00258932, 0.00258678, 0.00258754,
0.00259157, 0.00259742, 0.00260376, 0.00261015, 0.00261699, 0.00262534, 0.00263624,
0.00265, 0.00266603, 0.00268315, 0.00269997, 0.00271538, 0.0027293, 0.00274301,
0.00275855, 0.00277762, 0.00280069, 0.00282693, 0.00285474, 0.00288248, 0.00290876,
0.00293257, 0.00295365, 0.0029727, 0.00299094, 0.00300944, 0.0030288, 0.00304906,
0.00306975, 0.00309011, 0.00310917, 0.00312581, 0.00313895, 0.00314813, 0.00315378,
0.00315703, 0.0031592, 0.00316134, 0.00316407, 0.00316761, 0.00317182, 0.00317628,
0.00318021, 0.00318272, 0.00318301, 0.00318065, 0.00317563, 0.00316838, 0.00315962,
0.00315024, 0.00314115, 0.00313327, 0.00312741, 0.0031241, 0.00312351, 0.00312549,
0.00312965, 0.00313552, 0.00314259, 0.00315038, 0.00315846, 0.00316645, 0.00317408,
0.00318122, 0.00318795, 0.0031945, 0.00320115, 0.00320818, 0.0032158, 0.00322408,
0.00323293, 0.00324213, 0.00325147, 0.00326083, 0.00327028, 0.00328004, 0.00329046,
0.00330194, 0.00331489, 0.0033297, 0.00334668, 0.00336601, 0.00338771, 0.00341162,
0.0034374, 0.0034646, 0.00349265, 0.003521, 0.00354909, 0.00357647, 0.00360277,
0.00362773, 0.00365124, 0.00367333, 0.0036942, 0.00371413, 0.00373349, 0.00375265,
0.0037719, 0.00379146, 0.00381142, 0.00383177, 0.00385234, 0.0038729, 0.0038931,
0.00391258, 0.00393098, 0.00394801, 0.00396348, 0.00397735, 0.0039897, 0.00400073,
0.00401069, 0.0040199, 0.00402867, 0.00403726, 0.0040459, 0.00405473, 0.00406382,
0.00407317, 0.00408273, 0.00409241, 0.00410213, 0.00411181, 0.00412144, 0.00413098,
0.00414044, 0.00414982, 0.00415912, 0.00416834, 0.00417744, 0.00418636, 0.00419505,
0.00420339, 0.0042113, 0.00421866, 0.00422539, 0.00423144, 0.00423681, 0.00424153
];
let q = vec![0.85, 2.0];
let moving =
Transform::new(q.clone(), vec![fixed.shape()[0]]).interpolate::<3, _, _>(&fixed)?;
// The registration finds T such that im_b(T(x)) = im_a(x), which is the inverse of q
let q_inv = Transform::<ndarray::Ix1>::new(q.clone(), vec![fixed.shape()[0]]).inverse()?;
// Use default steps with ASGD — the multi-resolution pyramid converges
// reliably for this 1D affine problem.
let t = Transform::register(fixed.view(), moving.view(), vec![None; 2], None, None)?;
// Compare transformed coordinates of all pixels
let mut sum_diff = 0.0;
let mut count = 0;
for point in 0..fixed.len() {
let t_point = t.transform_point(&[point]);
let e_point = q_inv.transform_point(&[point]);
let diff_sq = (t_point[0] - e_point[0]).powi(2);
sum_diff += diff_sq.sqrt();
count += 1;
}
let mean_diff = sum_diff / count as f64;
println!("Our: {:?}\nmean_coord_diff: {:.6}", t, mean_diff);
println!("t: {:?}", t);
println!("i: {:?}", t.inverse()?);
println!("q_inv: {:?}", q_inv.parameters);
assert!(
t.parameters
.iter()
.zip(q_inv.parameters.iter())
.map(|(a, b)| (a - b).powi(2))
.sum::<f64>()
< 1.0
);
assert!(mean_diff < 0.1);
Ok(())
}
#[test]
fn diag_register1_metric() -> Result<(), Box<dyn std::error::Error>> {
use crate::bspline::{BSpline, BSplineTrait};
use crate::metric::{FixedMu, MattesMetric, SamplingArg, Sigma};
use algos::ObjectiveFunction;
let im_a = julia_image(
&[100, 1],
&[1.0, 0.0, 0.0, 0.01, 0.0, 0.0],
@@ -843,32 +933,92 @@ mod tests {
let im_b =
Transform::new(q.clone(), vec![im_a.shape()[0]]).interpolate::<1, _, _>(&im_a)?;
// The registration finds T such that im_b(T(x)) = im_a(x), which is the inverse of q
let f = Sigma::Absolute(vec![1.0]).smooth(im_a.view())?;
let m = Sigma::Absolute(vec![1.0]).smooth(im_b.view())?;
let metric = MattesMetric::<ndarray::Ix1>::new(
BSpline::new(f.view()),
BSpline::new(m.view()),
SamplingArg::Random(100),
32,
0.05,
)?
.with_fixed_mu(FixedMu::new_none(1));
let q_inv = Transform::<ndarray::Ix1>::new(q.clone(), vec![im_a.shape()[0]])
.inverse()?
.parameters;
// Use default steps with ASGD — the multi-resolution pyramid converges
// reliably for this 1D affine problem.
let t = Transform::register(im_a.view(), im_b.view(), vec![None; 2], None, None)?;
println!("t: {:?}", t);
println!("i: {:?}", t.inverse()?);
println!("q_inv: {:?}", q_inv);
assert!(
t.parameters
.iter()
.zip(q_inv.iter())
.map(|(a, b)| (a - b).powi(2))
.sum::<f64>()
< 1.0
// sweep scale at the optimal translation, and translation at scale 1.0
let mut best = (f64::MAX, vec![0.0, 0.0]);
let mut rows = Vec::new();
for s in (60..=360).map(|i| i as f64 / 100.0) {
for t in (-120..=80).map(|i| i as f64 / 40.0) {
let p = vec![s, t];
let v = metric.evaluate(&p);
if v < best.0 {
best = (v, p);
}
}
let v0 = metric.evaluate(&[s, -2.353]);
rows.push(format!("s={:.2} t=-2.353 val={:.4}", s, v0));
}
println!("BEST on grid: val={:.4} params={:?}", best.0, best.1);
println!("rows:");
// line from identity to q_inv and beyond
let mut prev = f64::NAN;
for (i, alpha) in (0..=40).map(|i| i as f64 / 10.0).enumerate() {
let p = vec![1.0 + (q_inv[0] - 1.0) * alpha, 0.0 + q_inv[1] * alpha];
let v = metric.evaluate(&p);
println!("line[{}] alpha={:.1} p={:?} val={:.4}", i, alpha, p, v);
if !prev.is_nan() && v > prev {
println!(" ^ INCREASING here — minimum is before this point");
}
prev = v;
}
// metric value along diagonal from q_inv to the registered optimum
for p in [q_inv.clone(), vec![2.0, -1.0], vec![2.5, -0.3]] {
println!("eval {:?}: val={:.4}", p, metric.evaluate(&p));
}
// numerical vs analytical gradient at identity and q_inv
let eps = 1e-5;
for p in [vec![1.0, 0.0], q_inv.clone()] {
let grad = metric.gradient(&p).unwrap();
let mut num = Vec::new();
for i in 0..2 {
let mut pp = p.clone();
let mut pm = p.clone();
pp[i] += eps;
pm[i] -= eps;
num.push((metric.evaluate(&pp) - metric.evaluate(&pm)) / (2.0 * eps));
}
println!("grad at {:?}: analytical={:?} numerical={:?}", p, grad, num);
}
let (t, results) = Transform::<ndarray::Ix1>::register_debug(
im_a.view(),
im_b.view(),
vec![None, None],
None,
None,
)?;
println!("register_debug t: {:?}", t);
for r in results {
println!(
" step: opt={:?} iters={} converged={} val={}",
r.optimal_point, r.iterations, r.converged, r.optimal_value
);
}
Ok(())
}
#[test]
fn register2_random_affine() -> Result<(), Box<dyn std::error::Error>> {
use rand::prelude::*;
let mut rng = rand::rngs::StdRng::seed_from_u64(1337);
let mut rng = StdRng::seed_from_u64(1337);
let shape = [200, 200];
let center = [99.5, 99.5];