diff --git a/Cargo.toml b/Cargo.toml index 9e44e70..1bffb81 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -31,6 +31,7 @@ tiffwrite = "2026.6.0" [dev-dependencies] tempfile = "3" +tiff = "0.11" [profile.release] debug = true diff --git a/src/metric.rs b/src/metric.rs index e1d3f5f..40ddb09 100644 --- a/src/metric.rs +++ b/src/metric.rs @@ -167,13 +167,19 @@ impl Sampling { fn index(&self) -> Vec> { match self { Self::Fixed(index) => index.clone(), - Self::Random((n_samples, shape, _cached)) => rand::rng() - .random_iter::() - .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(), + Self::Random((n_samples, shape, cached)) => { + let mut cache = cached.borrow_mut(); + if cache.is_empty() { + *cache = rand::rng() + .random_iter::() + .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() + } } } } diff --git a/src/register.rs b/src/register.rs index 7c7e0c1..2a1b4f3 100644 --- a/src/register.rs +++ b/src/register.rs @@ -59,35 +59,25 @@ impl RegistrationStep { /// Default registration steps matching elastix `FixedSmoothingImagePyramid`. /// - /// Elastix uses `MultiResolutionGaussianSmoothingPyramidImageFilter` which applies - /// Gaussian smoothing with σ = 0.5 × factor at **full resolution** — images are - /// **NOT downsampled**. See: `itkMultiResolutionGaussianSmoothingPyramidImageFilter.hxx` - /// - /// Schedule [8, 4, 2, 1] means σ = [4.0, 2.0, 1.0, 0.5] (spacing=1). - /// MaximumNumberOfIterations: 256 (matching SimpleElastix default affine). + /// Smooth-only multi-resolution pyramid (no downsampling). + /// Schedule [8, 4, 2, 1] → σ = [4.0, 2.0, 1.0, 0.5] at spacing=1. + /// Uses cached random sampling (NewSamplesEveryIteration=false per level). pub fn default_steps(ndim: usize, n: usize) -> Vec { let sigma_schedule: Vec = vec![4.0, 2.0, 1.0, 0.5]; - let nlevels = sigma_schedule.len(); - let mut steps = Vec::new(); - - for i in 0..nlevels { - 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_schedule[i]; ndim]), - SamplingArg::Random(samples), - 32, - 1e-6, - 0.05, - 256, - 1.0, - )); - } - - steps + let n_samples = (n / 16).clamp(2048, 8192); + sigma_schedule + .iter() + .map(|&s| Self { + sigma: Sigma::Absolute(vec![s; ndim]), + samples: SamplingArg::Random(n_samples), + n_bins: 32, + tolerance: 1e-6, + edge: 0.05, + max_iterations: 1500, + learning_rate: 4.0, + downsample: 1, + }) + .collect() } } diff --git a/src/transform.rs b/src/transform.rs index ed0864d..2966b25 100644 --- a/src/transform.rs +++ b/src/transform.rs @@ -198,14 +198,27 @@ impl Transform { /// find the transform which transforms moving into fixed, using fixed_mu to specify which /// parameters to keep fixed - pub fn register<'a, F, M, T, G>(fixed: F, moving: M, fixed_mu: G) -> Result + pub fn register<'a, F, M, T, G>( + fixed: F, + moving: M, + fixed_mu: G, + steps: Option>, + initial_guess: Option>, + ) -> Result where F: AsArray<'a, T, D>, M: AsArray<'a, T, D>, T: 'a + Clone + AsPrimitive, G: Into, { - Registration::new(fixed_mu).register(fixed, moving) + let mut registration = Registration::new(fixed_mu); + if let Some(steps) = steps { + registration.set_steps(steps); + } + if let Some(initial_guess) = initial_guess { + registration.set_initial_guess(initial_guess); + } + registration.register(fixed, moving) } /// find the transform which transforms moving into fixed and return the results of each @@ -504,8 +517,13 @@ mod tests { use crate::julia_image; use crate::transform::Transform; use itertools::Itertools; - use ndarray::{Ix2, s}; + use ndarray::{Array2, Ix2, s}; use num::traits::FloatConst; + use std::fs::File; + use std::path::Path; + use tiff::decoder::{Decoder, DecodingResult}; + use tiff::tags::Tag; + use tiffwrite::IJTiffFile; #[test] fn interpolate() -> Result<(), Box> { @@ -736,39 +754,9 @@ mod tests { .inverse()? .parameters; - let all_points: Vec> = (0..100).map(|j| vec![j as f64]).collect(); - let steps = vec![ - crate::register::RegistrationStep::new( - crate::metric::Sigma::Absolute(vec![8.0]), - crate::metric::SamplingArg::FixedAt(all_points.clone()), - 32, - 1e-4, - 0.05, - 200, - 1.0, - ), - crate::register::RegistrationStep::new( - crate::metric::Sigma::Absolute(vec![2.0]), - crate::metric::SamplingArg::FixedAt(all_points), - 64, - 1e-6, - 0.04, - 200, - 1.0, - ), - ]; - - let (t, steps) = Transform::register_debug( - im_a.view(), - im_b.view(), - vec![None, None], - Some(steps), - None, - )?; - println!("steps:"); - for step in steps { - println!(" {:?}", step); - } + // 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); @@ -1063,4 +1051,73 @@ mod tests { ); Ok(()) } + + fn read_tiff>(path: P) -> Result, Box> { + let mut reader = Decoder::new(File::open(path)?)?; + reader.seek_to_image(0)?; + let bytes = match reader.read_image()? { + DecodingResult::U8(data) => data.into_iter().map(|i| i as f64).collect::>(), + DecodingResult::U16(data) => data.into_iter().map(|i| i as f64).collect::>(), + DecodingResult::U32(data) => data.into_iter().map(|i| i as f64).collect::>(), + DecodingResult::U64(data) => data.into_iter().map(|i| i as f64).collect::>(), + DecodingResult::I8(data) => data.into_iter().map(|i| i as f64).collect::>(), + DecodingResult::I16(data) => data.into_iter().map(|i| i as f64).collect::>(), + DecodingResult::I32(data) => data.into_iter().map(|i| i as f64).collect::>(), + DecodingResult::I64(data) => data.into_iter().map(|i| i as f64).collect::>(), + DecodingResult::F16(data) => data.into_iter().map(|i| f64::from(i)).collect::>(), + DecodingResult::F32(data) => data.into_iter().map(|i| i as f64).collect::>(), + DecodingResult::F64(data) => data, + }; + let width = reader.get_tag(Tag::ImageWidth)?.into_u32()? as usize; + let height = reader.get_tag(Tag::ImageLength)?.into_u32()? as usize; + Ok(Array2::from_shape_vec((width, height), bytes)?) + } + + #[test] + fn register_real_images() -> Result<(), Box> { + let fixed = read_tiff("test_files/fixed.tif")?; + let moving = read_tiff("test_files/moving.tif")?; + + // SimpleElastix expected — the registration result is in TIFF-native coordinates. + let expected = [ + 0.9899559376493817, + 0.011269992506480442, + 0.017048860489651384, + 0.8533772512806085, + -12.877028775909979, + 4.143118928117275, + ]; + + let t = Transform::::register( + fixed.view(), + moving.view(), + vec![None, None, None, None, None, None], + None, + 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); + + assert!(max_err < 1.0); + assert!(sse < 1.0); + + let mut tif = IJTiffFile::new(std::env::home_dir().unwrap().join("tmp/register_real_images.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(moving.mapv(|i| i as u16), 2, 0, 0)?; + + Ok(()) + } } diff --git a/test_files/fixed.tif b/test_files/fixed.tif new file mode 100644 index 0000000..c97f3db Binary files /dev/null and b/test_files/fixed.tif differ diff --git a/test_files/moving.tif b/test_files/moving.tif new file mode 100644 index 0000000..7b2c827 Binary files /dev/null and b/test_files/moving.tif differ