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

- registration getting better

Browse Source
This commit is contained in:
Wim Pomp
2026-07-29 22:58:36 +02:00
parent 001df705ef
commit 36b24332cd
3 changed files with 114 additions and 82 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/metric.rs
+2 -2
View File
@@ -281,8 +281,8 @@ where
T: 'a + Clone + AsPrimitive<f64>,
{
Self::new(
BSpline::new(fixed),
BSpline::new(moving),
BSpline::<0, _>::new(fixed),
BSpline::<3, _>::new(moving),
sampling,
n_bins,
edge,
src/register.rs
+64 -48
View File
@@ -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<Self> {
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<D: Dimension> Registration<D> {
max_iterations,
learning_rate,
downsample: _,
optimizer,
} in steps.into_iter()
{
let f = sigma.smooth(fixed.view())?;
@@ -321,58 +326,68 @@ impl<D: Dimension> Registration<D> {
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::<D>::new(p, fixed.shape().to_vec()))
}
fn optimize_metric<M: ObjectiveFunction<f64>>(
metric: &M,
p_var: &[f64],
optimizer: &Optimizer,
max_iterations: usize,
tolerance: f64,
learning_rate: f64,
scales: &[f64],
) -> Option<algos::OptimizationResult<f64>> {
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<D: Dimension> Registration<D> {
max_iterations,
learning_rate,
downsample: _,
optimizer,
} in steps.into_iter()
{
let f = sigma.smooth(fixed.view())?;
@@ -434,7 +450,7 @@ impl<D: Dimension> Registration<D> {
// 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,
src/transform.rs
+48 -32
View File
@@ -1008,20 +1008,24 @@ mod tests {
None,
)?;
let sse = t
.parameters
.iter()
.zip(expected.iter())
.map(|(a, b)| (a - b).powi(2))
.sum::<f64>();
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::<Ix2>::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::<f64>();
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(())
}