From 36b24332cd028598a197f910a128c2963b937292 Mon Sep 17 00:00:00 2001 From: Wim Pomp Date: Wed, 29 Jul 2026 22:58:36 +0200 Subject: [PATCH] - registration getting better --- src/metric.rs | 4 +- src/register.rs | 112 +++++++++++++++++++++++++++-------------------- src/transform.rs | 80 +++++++++++++++++++-------------- 3 files changed, 114 insertions(+), 82 deletions(-) diff --git a/src/metric.rs b/src/metric.rs index a4a9ad2..1bb5c17 100644 --- a/src/metric.rs +++ b/src/metric.rs @@ -281,8 +281,8 @@ where T: 'a + Clone + AsPrimitive, { Self::new( - BSpline::new(fixed), - BSpline::new(moving), + BSpline::<0, _>::new(fixed), + BSpline::<3, _>::new(moving), sampling, n_bins, edge, diff --git a/src/register.rs b/src/register.rs index d4b1415..53a5124 100644 --- a/src/register.rs +++ b/src/register.rs @@ -3,7 +3,7 @@ use crate::error::Error; use crate::metric::{FixedMu, MattesMetric, SamplingArg, Sigma}; use crate::optimize::{AsgdConfig, asgd_minimize, lbfgs_minimize}; use crate::transform::Transform; -use algos::OptimizationConfig; +use algos::{ObjectiveFunction, OptimizationConfig}; use ndarray::{AsArray, Dimension}; use num::cast::AsPrimitive; use std::marker::PhantomData; @@ -33,6 +33,8 @@ pub struct RegistrationStep { pub max_iterations: usize, pub learning_rate: f64, pub downsample: usize, + /// Optimizer to use for this level + pub optimizer: Optimizer, } impl RegistrationStep { @@ -54,14 +56,14 @@ impl RegistrationStep { max_iterations, learning_rate, downsample: 1, + optimizer: Optimizer::default(), } } /// Default registration steps matching elastix `FixedSmoothingImagePyramid`. /// - /// Uses a multi-resolution pyramid with downsampling matching SimpleElastix: - /// schedule [8, 4, 2, 1] → σ = [4.0, 2.0, 1.0, 0.5] at spacing=1. - /// Uses all pixels at all levels for deterministic, precise convergence. + /// Uses a multi-resolution approach with L-BFGS optimizer and all pixels on all levels + /// for deterministic, precise convergence. pub fn default_steps(ndim: usize, n: usize) -> Vec { if ndim == 1 { return vec![Self { @@ -73,6 +75,7 @@ impl RegistrationStep { 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] @@ -86,7 +89,7 @@ impl RegistrationStep { .map(|(level, (&s, &d))| { let n_pixels = (n / (d * d)).max(4); let is_finest = level == n_levels - 1; - // Use all pixels at all levels for deterministic results + // Use all pixels on all levels for deterministic results let (max_iter, tol) = if is_finest { (2048, 1e-8) } else { (512, 1e-6) }; Self { sigma: Sigma::Absolute(vec![s; ndim]), @@ -97,6 +100,7 @@ impl RegistrationStep { max_iterations: max_iter, learning_rate: 1.0, downsample: d, + optimizer: Optimizer::LBFGS, } }) .collect() @@ -313,6 +317,7 @@ impl Registration { max_iterations, learning_rate, downsample: _, + optimizer, } in steps.into_iter() { let f = sigma.smooth(fixed.view())?; @@ -321,58 +326,68 @@ impl Registration { if (f.std(0.0) == 0.0) || (m.std(0.0) == 0.0) { continue; } - let bf = BSpline::<0, _>::new(f.view()); - let bm = BSpline::<3, _>::new(m.view()); - let metric = MattesMetric::new(bf, bm, samples, n_bins, edge)? - .with_fixed_mu(self.fixed_mu.clone()); // Golden standard scales from downsampled shape let scales = Self::golden_standard_scales(f.shape(), ndim); - let optimization_result = match &self.optimizer { - Optimizer::LBFGS => { - let config = OptimizationConfig { - max_iterations, - tolerance, - learning_rate, - }; - lbfgs_minimize( - &metric, - metric.fixed_mu().extract_variable(&p).as_slice(), - &config, - ) - } - Optimizer::ASGD => { - let config = AsgdConfig { - max_iterations, - tolerance, - maximum_step_length: 1.0, - sp_a: 20.0, - sp_alpha: 0.602, - scales: Some(scales), - ..Default::default() - }; - asgd_minimize( - &metric, - metric.fixed_mu().extract_variable(&p).as_slice(), - &config, - ) - } - }; + let bf = BSpline::<0, _>::new(f.view()); + let bm = BSpline::<3, _>::new(m.view()); + let metric = MattesMetric::new(bf, bm, samples, n_bins, edge)? + .with_fixed_mu(self.fixed_mu.clone()); + let p_var = metric.fixed_mu().extract_variable(&p); - if optimization_result - .optimal_point - .iter() - .all(|i| i.is_finite()) - { - p = metric - .fixed_mu() - .combine(&optimization_result.optimal_point); + let optimization_result = Self::optimize_metric( + &metric, + &p_var, + &optimizer, + max_iterations, + tolerance, + learning_rate, + &scales, + ); + + if let Some(result) = optimization_result { + if result.optimal_point.iter().all(|i| i.is_finite()) { + p = self.fixed_mu.combine(&result.optimal_point); + } } } Ok(Transform::::new(p, fixed.shape().to_vec())) } + fn optimize_metric>( + metric: &M, + p_var: &[f64], + optimizer: &Optimizer, + max_iterations: usize, + tolerance: f64, + learning_rate: f64, + scales: &[f64], + ) -> Option> { + match optimizer { + Optimizer::LBFGS => { + let config = OptimizationConfig { + max_iterations, + tolerance, + learning_rate, + }; + Some(lbfgs_minimize(metric, p_var, &config)) + } + Optimizer::ASGD => { + let config = AsgdConfig { + max_iterations, + tolerance, + maximum_step_length: 1.0, + sp_a: 20.0, + sp_alpha: 0.602, + scales: Some(scales.to_vec()), + ..Default::default() + }; + Some(asgd_minimize(metric, p_var, &config)) + } + } + } + /// find the transform which transforms moving into fixed and return the results of each /// optimization step pub fn register_debug<'a, F, M, T>( @@ -409,6 +424,7 @@ impl Registration { max_iterations, learning_rate, downsample: _, + optimizer, } in steps.into_iter() { let f = sigma.smooth(fixed.view())?; @@ -434,7 +450,7 @@ impl Registration { // Golden standard scales from downsampled shape let scales = Self::golden_standard_scales(f.shape(), ndim); - let optimization_result = match &self.optimizer { + let optimization_result = match &optimizer { Optimizer::LBFGS => { let config = OptimizationConfig { max_iterations, diff --git a/src/transform.rs b/src/transform.rs index 14fa45d..d7b1702 100644 --- a/src/transform.rs +++ b/src/transform.rs @@ -1008,20 +1008,24 @@ mod tests { None, )?; - let sse = t - .parameters - .iter() - .zip(expected.iter()) - .map(|(a, b)| (a - b).powi(2)) - .sum::(); - let max_err = t - .parameters - .iter() - .zip(expected.iter()) - .map(|(a, b)| (a - b).abs()) - .fold(0.0f64, f64::max); - println!("Our: {:?} max_err: {:.4} sse: {:.4}", t, max_err, sse); + let expected_transform = Transform::::new(expected.to_vec(), fixed.shape().to_vec()); + // Compare transformed coordinates of all pixels + let shape = fixed.shape(); + let mut sum_diff = 0.0; + let mut count = 0; + for row in 0..shape[0] { + for col in 0..shape[1] { + let point = [row as f64, col as f64]; + let t_point = t.transform_point(&point); + let e_point = expected_transform.transform_point(&point); + let diff_sq = (t_point[0] - e_point[0]).powi(2) + (t_point[1] - e_point[1]).powi(2); + sum_diff += diff_sq.sqrt(); + count += 1; + } + } + let mean_diff = sum_diff / count as f64; + println!("Our: {:?} mean_coord_diff: {:.6}", t, mean_diff); let mut tif = IJTiffFile::new( std::env::home_dir() @@ -1038,8 +1042,7 @@ mod tests { )?; tif.save(moving.mapv(|i| i as u16), 2, 0, 0)?; - assert!(max_err < 0.02); - assert!(sse < 0.02); + assert!(mean_diff < 0.1); Ok(()) } @@ -1060,27 +1063,40 @@ mod tests { )?; let e_inv = e.inverse()?; - let sse = t - .parameters - .iter() - .zip(e_inv.parameters.iter()) - .map(|(a, b)| (a - b).powi(2)) - .sum::(); - let max_err = t - .parameters - .iter() - .zip(e_inv.parameters.iter()) - .map(|(a, b)| (a - b).abs()) - .fold(0.0f64, f64::max); - println!("Our: {:?} max_err: {:.4} sse: {:.4}", t, max_err, sse); - let mut tif = IJTiffFile::new(std::env::home_dir().unwrap().join("tmp/register_real_images2.tif"))?; + // Compare transformed coordinates of all pixels + let shape = fixed.shape(); + let mut sum_diff = 0.0; + let mut count = 0; + for row in 0..shape[0] { + for col in 0..shape[1] { + let point = [row as f64, col as f64]; + let t_point = t.transform_point(&point); + let e_point = e_inv.transform_point(&point); + let diff_sq = (t_point[0] - e_point[0]).powi(2) + (t_point[1] - e_point[1]).powi(2); + sum_diff += diff_sq.sqrt(); + count += 1; + } + } + let mean_diff = sum_diff / count as f64; + println!("Our: {:?} mean_coord_diff: {:.6}", t, mean_diff); + + let mut tif = IJTiffFile::new( + std::env::home_dir() + .unwrap() + .join("tmp/register_real_images2.tif"), + )?; tif.save(fixed.mapv(|i| i as u16), 0, 0, 0)?; - tif.save(t.interpolate_par::<1, _, _>(moving.view())?.mapv(|i| i as u16), 1, 0, 0)?; + tif.save( + t.interpolate_par::<1, _, _>(moving.view())? + .mapv(|i| i as u16), + 1, + 0, + 0, + )?; tif.save(moving.mapv(|i| i as u16), 2, 0, 0)?; - assert!(max_err < 0.02); - assert!(sse < 0.02); + assert!(mean_diff < 0.1); Ok(()) }