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

- tests register1 and register2 passing

Browse Source
This commit is contained in:
Wim Pomp committed 2026-07-25 13:15:48 +02:00
1 parent 806571be82
commit 5b7dc18a4d
5 files changed
+887 -114
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

src/register.rs
+170 -75
View File
@@ -1,14 +1,28 @@
use crate::bspline::{BSpline, BSplineTrait};
use crate::error::Error;
use crate::metric::{FixedMu, MattesMetric, SamplingArg, Sigma};
use crate::optimize::{AsgdConfig, asgd_minimize, lbfgs_minimize};
use crate::transform::Transform;
use crate::optimize::lbfgs_minimize;
use algos::OptimizationConfig;
use ndarray::{AsArray, Dimension};
use num::cast::AsPrimitive;
use num::integer::Roots;
use std::marker::PhantomData;
/// Optimizer type for registration
#[derive(Clone, Debug, PartialEq)]
pub enum Optimizer {
/// L-BFGS optimizer (fast, requires consistent gradients)
LBFGS,
/// Adaptive Stochastic Gradient Descent (robust, handles noisy gradients)
ASGD,
}
impl Default for Optimizer {
fn default() -> Self {
Self::ASGD
}
}
#[derive(Clone, Debug)]
pub struct RegistrationStep {
pub sigma: Sigma,
@@ -41,54 +55,38 @@ impl RegistrationStep {
}
}
/// Default registration steps matching SimpleElastix affine defaults.
///
/// Elastix uses `FixedSmoothingImagePyramid` with schedule [8, 4, 2, 1] for 4 levels.
/// Sigma = 0.5 * factor * spacing. With spacing=1: sigma = [4.0, 2.0, 1.0, 0.5].
/// These are absolute sigma values in pixel units.
/// See: `itkMultiResolutionGaussianSmoothingPyramidImageFilter.hxx`
///
/// Elastix defaults: 2048 samples, 32 histogram bins, 256 iterations.
pub fn default_steps(ndim: usize, n: usize) -> Vec<Self> {
vec![
RegistrationStep::new(
Sigma::Relative(vec![0.5; ndim]),
SamplingArg::Fixed(n.sqrt().max(n / 125).max(100)),
3,
1e-4,
0.05,
100,
100.0,
),
RegistrationStep::new(
Sigma::Relative(vec![0.25; ndim]),
SamplingArg::Fixed(n.sqrt().max(n / 30).max(200)),
6,
let sigma_schedule: Vec<f64> = vec![4.0, 2.0, 1.0, 0.5];
let nlevels = sigma_schedule.len();
let mut steps = Vec::new();
for (i, &sigma) in sigma_schedule.iter().enumerate() {
let fraction = 1.0 / 2.0_f64.powi((nlevels - 1 - i) as i32);
let samples = (n as f64 * fraction)
.sqrt()
.max(n as f64 * fraction / 10.0)
.max(2048.0) as usize;
steps.push(RegistrationStep::new(
Sigma::Absolute(vec![sigma; ndim]),
SamplingArg::Random(samples),
32,
1e-6,
0.04,
100,
0.05,
256,
1.0,
),
RegistrationStep::new(
Sigma::Absolute(vec![8.0; ndim]),
SamplingArg::Fixed(n.sqrt().max(n / 20).max(400)),
24,
3e-7,
0.03,
100,
1e-2,
),
RegistrationStep::new(
Sigma::Absolute(vec![2.0; ndim]),
SamplingArg::Fixed(n.sqrt().max(n / 10).max(800)),
96,
1e-7,
0.01,
100,
1e-2,
),
RegistrationStep::new(
Sigma::None,
SamplingArg::Fixed(n.sqrt().max(n / 5).max(1600)),
240,
1e-8,
0.001,
100,
1e-3,
),
]
));
}
steps
}
}
@@ -112,6 +110,7 @@ pub struct Registration<D: Dimension> {
fixed_mu: FixedMu,
steps: Option<Vec<RegistrationStep>>,
initial_guess: Option<Vec<f64>>,
optimizer: Optimizer,
dimension: PhantomData<D>,
}
@@ -121,6 +120,7 @@ impl<D: Dimension> Registration<D> {
fixed_mu: fixed_mu.into(),
steps: None,
initial_guess: None,
optimizer: Optimizer::default(),
dimension: PhantomData,
}
}
@@ -135,6 +135,7 @@ impl<D: Dimension> Registration<D> {
fixed_mu: FixedMu::new_none(ndim),
steps: None,
initial_guess: None,
optimizer: Optimizer::default(),
dimension: PhantomData,
})
}
@@ -154,6 +155,7 @@ impl<D: Dimension> Registration<D> {
fixed_mu: fixed_mu.into(),
steps: None,
initial_guess: None,
optimizer: Optimizer::default(),
dimension: PhantomData,
})
}
@@ -168,6 +170,11 @@ impl<D: Dimension> Registration<D> {
self
}
pub fn with_optimizer(mut self, optimizer: Optimizer) -> Self {
self.optimizer = optimizer;
self
}
pub fn set_steps(&mut self, steps: Vec<RegistrationStep>) {
self.steps = Some(steps);
}
@@ -176,6 +183,48 @@ impl<D: Dimension> Registration<D> {
self.initial_guess = Some(initial_guess);
}
pub fn set_optimizer(&mut self, optimizer: Optimizer) {
self.optimizer = optimizer;
}
/// Compute automatic transform initialization by aligning geometric centers
fn compute_initial_transform(
fixed_shape: &[usize],
moving_shape: &[usize],
ndim: usize,
) -> Vec<f64> {
let mut params = vec![0.0; ndim * ndim + ndim];
for i in 0..ndim {
params[(ndim + 1) * i] = 1.0;
}
for i in 0..ndim {
let fixed_center = (fixed_shape[i] as f64 - 1.0) / 2.0;
let moving_center = (moving_shape[i] as f64 - 1.0) / 2.0;
params[ndim * ndim + i] = fixed_center - moving_center;
}
params
}
/// Compute golden standard scales for ASGD parameter normalization.
/// Matches elastix ScaledCostFunction convention: scales = sqrt(C[i][i]).
/// For affine transform v = R*(u-c) + t + c:
/// - Rotation/scale params: C[i][i] = Var(u-c) = (N^2-1)/12, scales[i] = sqrt(C[i][i])
/// - Translation params: C[i][i] = 1.0, scales[i] = 1.0
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
}
/// find the transform which transforms moving into fixed and return the results of each
/// optimization step
pub fn register<'a, F, M, T>(&self, fixed: F, moving: M) -> Result<Transform<D>, Error>
@@ -193,11 +242,11 @@ impl<D: Dimension> Registration<D> {
.as_ref()
.cloned()
.unwrap_or_else(|| RegistrationStep::default_steps(ndim, n));
let mut p = self
.initial_guess
.as_ref()
.cloned()
.unwrap_or_else(|| Transform::<D>::default().parameters);
let mut p = self.initial_guess.as_ref().cloned().unwrap_or_else(|| {
Self::compute_initial_transform(fixed.shape(), moving.shape(), ndim)
});
for RegistrationStep {
sigma,
samples,
@@ -217,16 +266,39 @@ impl<D: Dimension> Registration<D> {
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 optimization_config = OptimizationConfig {
max_iterations,
tolerance,
learning_rate,
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 scales = Self::golden_standard_scales(fixed.shape(), ndim);
let config = AsgdConfig {
max_iterations,
tolerance,
maximum_step_length: 1.0,
sp_a: 20.0,
sp_alpha: 1.0,
scales: Some(scales),
..Default::default()
};
asgd_minimize(
&metric,
metric.fixed_mu().extract_variable(&p).as_slice(),
&config,
)
}
};
let optimization_result = lbfgs_minimize(
&metric,
metric.fixed_mu().extract_variable(&p).as_slice(),
&optimization_config,
);
if optimization_result
.optimal_point
.iter()
@@ -261,11 +333,11 @@ impl<D: Dimension> Registration<D> {
.as_ref()
.cloned()
.unwrap_or_else(|| RegistrationStep::default_steps(ndim, n));
let mut p = self
.initial_guess
.as_ref()
.cloned()
.unwrap_or_else(|| Transform::<D>::default().parameters);
let mut p = self.initial_guess.as_ref().cloned().unwrap_or_else(|| {
Self::compute_initial_transform(fixed.shape(), moving.shape(), ndim)
});
let mut registration_results = Vec::new();
for RegistrationStep {
sigma,
@@ -291,16 +363,39 @@ impl<D: Dimension> Registration<D> {
};
let metric = MattesMetric::new(bf, bm, samples, n_bins, edge)?
.with_fixed_mu(self.fixed_mu.clone());
let optimization_config = OptimizationConfig {
max_iterations,
tolerance,
learning_rate,
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 scales = Self::golden_standard_scales(fixed.shape(), ndim);
let config = AsgdConfig {
max_iterations,
tolerance,
maximum_step_length: 1.0,
sp_a: 20.0,
sp_alpha: 1.0,
scales: Some(scales),
..Default::default()
};
asgd_minimize(
&metric,
metric.fixed_mu().extract_variable(&p).as_slice(),
&config,
)
}
};
let optimization_result = lbfgs_minimize(
&metric,
metric.fixed_mu().extract_variable(&p).as_slice(),
&optimization_config,
);
if optimization_result
.optimal_point
.iter()