1104 lines
36 KiB
Rust
1104 lines
36 KiB
Rust
use crate::bspline::{BSpline, BSplineTrait};
|
||
use crate::error::Error;
|
||
use crate::metric::FixedMu;
|
||
use crate::register::{Registration, RegistrationResult, RegistrationStep};
|
||
use itertools::Itertools;
|
||
use ndarray::{Array, Array2, ArrayD, AsArray, Dimension, Ix2, s};
|
||
use ndarray_linalg::Inverse;
|
||
use num::cast::AsPrimitive;
|
||
use serde::{Deserialize, Serialize};
|
||
use serde_yaml::{from_reader, to_writer};
|
||
use std::fs::File;
|
||
use std::marker::PhantomData;
|
||
use std::ops::Mul;
|
||
use std::path::PathBuf;
|
||
|
||
/// get coordinates resulting from transforming input coordinates, coordinate must have N
|
||
/// columns: x, y, z, ...
|
||
#[inline]
|
||
pub fn transform_point<'a, T>(point: &'a [T], center: &'a [f64], parameters: &'a [f64]) -> Vec<f64>
|
||
where
|
||
T: AsPrimitive<f64>,
|
||
{
|
||
// [a, b, c] x = a*x + b*y + c
|
||
// [d, e, f] y = d*x + e*y + f
|
||
// [0, 0, 1] 1 = 1
|
||
debug_assert_eq!(center.len(), point.len());
|
||
debug_assert_eq!((center.len() + 1) * center.len(), parameters.len());
|
||
|
||
let m = center.len();
|
||
let n = m * m;
|
||
let point = point
|
||
.iter()
|
||
.zip_eq(center.iter())
|
||
.map(|(p, c)| p.as_() - c)
|
||
.collect::<Vec<_>>();
|
||
parameters
|
||
.iter()
|
||
.take(n)
|
||
.chunks(m)
|
||
.into_iter()
|
||
.zip_eq(parameters.iter().skip(n))
|
||
.zip_eq(center.iter())
|
||
.map(|((r, t), c)| r.zip_eq(point.iter()).map(|(i, p)| i * p).sum::<f64>() + t + c)
|
||
.collect()
|
||
}
|
||
|
||
/// a struct describing the transform
|
||
/// generic parameter N = # image dimensions
|
||
#[derive(Clone, Debug, Deserialize, Serialize, PartialEq)]
|
||
pub struct Transform<D: Dimension = Ix2> {
|
||
/// flattened NxN rotation matrix + N translation parameters
|
||
pub parameters: Vec<f64>,
|
||
/// error / significance on parameters
|
||
pub dparameters: Vec<f64>,
|
||
/// the point about which rotations are performed
|
||
pub center: Vec<f64>,
|
||
/// the shape of images for which this transform is meant
|
||
pub shape: Vec<usize>,
|
||
ndim: usize,
|
||
dimension: PhantomData<D>,
|
||
}
|
||
|
||
impl<D: Dimension> Mul for Transform<D> {
|
||
type Output = Transform<D>;
|
||
|
||
fn mul(self, rhs: Self) -> Self::Output {
|
||
&self * &rhs
|
||
}
|
||
}
|
||
|
||
impl<D: Dimension> Mul<&Transform<D>> for Transform<D> {
|
||
type Output = Transform<D>;
|
||
|
||
fn mul(self, rhs: &Transform<D>) -> Self::Output {
|
||
&self * rhs
|
||
}
|
||
}
|
||
|
||
impl<D: Dimension> Mul<Transform<D>> for &Transform<D> {
|
||
type Output = Transform<D>;
|
||
|
||
fn mul(self, rhs: Transform<D>) -> Self::Output {
|
||
self * &rhs
|
||
}
|
||
}
|
||
|
||
impl<D: Dimension> Mul<&Transform<D>> for &Transform<D> {
|
||
type Output = Transform<D>;
|
||
|
||
#[allow(clippy::suspicious_arithmetic_impl)]
|
||
fn mul(self, rhs: &Transform<D>) -> Self::Output {
|
||
let m = self.matrix().dot(&rhs.matrix());
|
||
let dm = self.dmatrix().dot(&rhs.matrix()) + self.matrix().dot(&rhs.dmatrix());
|
||
Self::Output {
|
||
parameters: m
|
||
.slice(s![..self.ndim, ..self.ndim])
|
||
.flatten()
|
||
.iter()
|
||
.chain(m.slice(s![..self.ndim, self.ndim]).iter())
|
||
.cloned()
|
||
.collect(),
|
||
dparameters: dm
|
||
.slice(s![..self.ndim, ..])
|
||
.flatten()
|
||
.iter()
|
||
.chain(dm.slice(s![..self.ndim, ..]).iter())
|
||
.cloned()
|
||
.collect(),
|
||
center: self.center.clone(),
|
||
shape: self.shape.clone(),
|
||
ndim: self.ndim,
|
||
dimension: self.dimension,
|
||
}
|
||
}
|
||
}
|
||
|
||
impl<D: Dimension> Eq for Transform<D> {}
|
||
|
||
impl<D: Dimension> Default for Transform<D> {
|
||
/// the unit transform
|
||
fn default() -> Self {
|
||
let ndim = D::NDIM.expect("number of dimensions must be known to initialise transform");
|
||
let mut parameters = vec![0.0; ndim * ndim + ndim];
|
||
for i in 0..ndim {
|
||
parameters[(ndim + 1) * i] = 1.0;
|
||
}
|
||
|
||
Self {
|
||
parameters,
|
||
dparameters: vec![0f64; ndim * ndim + ndim],
|
||
center: vec![0f64; ndim],
|
||
shape: vec![0usize; ndim],
|
||
ndim,
|
||
dimension: PhantomData,
|
||
}
|
||
}
|
||
}
|
||
|
||
impl<D: Dimension> Transform<D> {
|
||
/// parameters: flat NxN part of matrix + translation; center: center of rotation
|
||
pub fn new_with_center(parameters: Vec<f64>, center: Vec<f64>, shape: Vec<usize>) -> Self {
|
||
let ndim = if let Some(ndim) = D::NDIM {
|
||
ndim
|
||
} else {
|
||
center.len()
|
||
};
|
||
debug_assert_eq!(parameters.len(), ndim * ndim + ndim);
|
||
debug_assert_eq!(center.len(), ndim);
|
||
debug_assert_eq!(shape.len(), ndim);
|
||
Self {
|
||
parameters,
|
||
dparameters: vec![0.0; ndim * ndim + ndim],
|
||
center,
|
||
shape,
|
||
ndim,
|
||
dimension: PhantomData,
|
||
}
|
||
}
|
||
|
||
pub fn new(parameters: Vec<f64>, shape: Vec<usize>) -> Self {
|
||
let ndim = if let Some(ndim) = D::NDIM {
|
||
ndim
|
||
} else {
|
||
shape.len()
|
||
};
|
||
debug_assert_eq!(parameters.len(), ndim * ndim + ndim);
|
||
debug_assert_eq!(shape.len(), ndim);
|
||
let center = shape.iter().map(|s| ((*s as f64) - 1.0) / 2.0).collect();
|
||
Self {
|
||
parameters,
|
||
dparameters: vec![0.0; ndim * ndim + ndim],
|
||
center,
|
||
shape,
|
||
ndim,
|
||
dimension: PhantomData,
|
||
}
|
||
}
|
||
|
||
/// find the affine transform which transforms moving into fixed
|
||
pub fn register_affine<'a, F, M, T>(fixed: F, moving: M) -> Result<Self, Error>
|
||
where
|
||
F: AsArray<'a, T, D>,
|
||
M: AsArray<'a, T, D>,
|
||
T: 'a + Clone + AsPrimitive<f64>,
|
||
{
|
||
Registration::new_affine()?.register(fixed, moving)
|
||
}
|
||
|
||
/// find the translation which transforms moving into fixed
|
||
pub fn register_translation<'a, F, M, T>(fixed: F, moving: M) -> Result<Self, Error>
|
||
where
|
||
F: AsArray<'a, T, D>,
|
||
M: AsArray<'a, T, D>,
|
||
T: 'a + Clone + AsPrimitive<f64>,
|
||
{
|
||
Registration::new_translation()?.register(fixed, moving)
|
||
}
|
||
|
||
/// 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,
|
||
steps: Option<Vec<RegistrationStep>>,
|
||
initial_guess: Option<Vec<f64>>,
|
||
) -> Result<Self, Error>
|
||
where
|
||
F: AsArray<'a, T, D>,
|
||
M: AsArray<'a, T, D>,
|
||
T: 'a + Clone + AsPrimitive<f64>,
|
||
G: Into<FixedMu>,
|
||
{
|
||
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
|
||
/// optimization step
|
||
pub fn register_debug<'a, F, M, T, G>(
|
||
fixed: F,
|
||
moving: M,
|
||
fixed_mu: G,
|
||
steps: Option<Vec<RegistrationStep>>,
|
||
initial_guess: Option<Vec<f64>>,
|
||
) -> Result<(Self, Vec<RegistrationResult>), Error>
|
||
where
|
||
F: AsArray<'a, T, D>,
|
||
M: AsArray<'a, T, D>,
|
||
T: 'a + Clone + AsPrimitive<f64>,
|
||
G: Into<FixedMu>,
|
||
{
|
||
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_debug(fixed, moving)
|
||
}
|
||
|
||
/// create a transform from a translation
|
||
pub fn from_translation(translation: &[f64]) -> Self {
|
||
let ndim = if let Some(ndim) = D::NDIM {
|
||
ndim
|
||
} else {
|
||
translation.len()
|
||
};
|
||
debug_assert_eq!(translation.len(), ndim);
|
||
let mut parameters = vec![0.0; ndim * ndim + ndim];
|
||
for i in 0..ndim {
|
||
parameters[(ndim + 1) * i] = 1.0;
|
||
}
|
||
for (p, t) in parameters
|
||
.iter_mut()
|
||
.skip(ndim * ndim)
|
||
.zip_eq(translation.iter())
|
||
{
|
||
*p = *t;
|
||
}
|
||
|
||
Self {
|
||
parameters,
|
||
dparameters: vec![0f64; ndim * ndim + ndim],
|
||
center: vec![0f64; ndim],
|
||
shape: vec![0usize; ndim],
|
||
ndim,
|
||
dimension: PhantomData,
|
||
}
|
||
}
|
||
|
||
pub fn with_shape(mut self, shape: Vec<usize>) -> Self {
|
||
self.shape = shape;
|
||
self
|
||
}
|
||
|
||
/// create a transform from a scaling
|
||
pub fn from_scaling(scaling: &[f64]) -> Self {
|
||
let ndim = if let Some(ndim) = D::NDIM {
|
||
ndim
|
||
} else {
|
||
scaling.len()
|
||
};
|
||
debug_assert_eq!(scaling.len(), ndim);
|
||
let mut parameters = vec![0.0; ndim * ndim + ndim];
|
||
parameters[ndim * ndim + ndim - 1] = 1.0;
|
||
for (p, s) in parameters
|
||
.iter_mut()
|
||
.step_by(ndim + 1)
|
||
.zip_eq(scaling.iter())
|
||
{
|
||
*p = *s;
|
||
}
|
||
|
||
Self {
|
||
parameters,
|
||
dparameters: vec![0f64; ndim * ndim + ndim],
|
||
center: vec![0f64; ndim],
|
||
shape: vec![0usize; ndim],
|
||
ndim,
|
||
dimension: PhantomData,
|
||
}
|
||
}
|
||
|
||
/// add a translation to self
|
||
pub fn with_translation(&self, translation: &[f64]) -> Self {
|
||
Self::from_translation(translation) * self
|
||
}
|
||
|
||
/// scale self
|
||
pub fn with_scaling(&self, scaling: &[f64]) -> Self {
|
||
Self::from_scaling(scaling) * self
|
||
}
|
||
|
||
/// read a transform from a file
|
||
pub fn from_file(path: PathBuf) -> Result<Self, Error> {
|
||
let file = File::open(path)?;
|
||
Ok(from_reader(file)?)
|
||
}
|
||
|
||
/// write a transform to a file
|
||
pub fn to_file(&self, path: PathBuf) -> Result<(), Error> {
|
||
let mut file = std::fs::OpenOptions::new()
|
||
.create(true)
|
||
.write(true)
|
||
.truncate(true)
|
||
.open(path)?;
|
||
to_writer(&mut file, self)?;
|
||
Ok(())
|
||
}
|
||
|
||
/// true if transform does nothing
|
||
pub fn is_unity(&self) -> bool {
|
||
let n = self.ndim * self.ndim;
|
||
let m = self.ndim + 1;
|
||
self.parameters
|
||
.iter()
|
||
.take(n)
|
||
.enumerate()
|
||
.all(|(i, &x)| if i % m == 0 { x == 1.0 } else { x == 0.0 })
|
||
&& self.parameters.iter().skip(n).all(|&x| x == 0.0)
|
||
}
|
||
|
||
/// get coordinates resulting from transforming input coordinates, coordinate must have N
|
||
/// columns: x, y, z, ...
|
||
pub fn transform_point<'a, T>(&'a self, point: &'a [T]) -> Vec<f64>
|
||
where
|
||
T: AsPrimitive<f64>,
|
||
{
|
||
transform_point(point, &self.center, &self.parameters)
|
||
}
|
||
|
||
/// get coordinates resulting from transforming input coordinates, coordinates must have N
|
||
/// columns: x, y, z, ...
|
||
pub fn transform_points<'a, A, T>(&'a self, points: A) -> Result<Array2<f64>, Error>
|
||
where
|
||
T: AsPrimitive<f64>,
|
||
A: AsArray<'a, T, Ix2>,
|
||
{
|
||
// [a, b, c] x = a*x + b*y + c
|
||
// [d, e, f] y = d*x + e*y + f
|
||
// [0, 0, 1] 1 = 1
|
||
let points = points.into();
|
||
let shape = points.shape();
|
||
let a = points
|
||
.rows()
|
||
.into_iter()
|
||
.flat_map(|p| {
|
||
if let Some(q) = p.as_slice() {
|
||
transform_point(q, &self.center, &self.parameters)
|
||
} else {
|
||
transform_point(&p.to_vec(), &self.center, &self.parameters)
|
||
}
|
||
})
|
||
.collect::<Vec<_>>();
|
||
Ok(ArrayD::from_shape_vec(shape, a)?.into_dimensionality()?)
|
||
}
|
||
|
||
/// get the matrix defining the transform
|
||
pub fn matrix(&self) -> Array2<f64> {
|
||
let n = self.ndim * self.ndim;
|
||
let mut matrix = Array2::eye(self.ndim + 1);
|
||
for (m, p) in matrix
|
||
.slice_mut(s![..self.ndim, ..self.ndim])
|
||
.iter_mut()
|
||
.zip_eq(self.parameters.iter().take(n))
|
||
{
|
||
*m = *p;
|
||
}
|
||
for (m, p) in matrix
|
||
.slice_mut(s![..self.ndim, self.ndim])
|
||
.iter_mut()
|
||
.zip_eq(self.parameters.iter().skip(n))
|
||
{
|
||
*m = *p;
|
||
}
|
||
matrix
|
||
}
|
||
|
||
/// get the matrix describing the error of the transform
|
||
pub fn dmatrix(&self) -> Array2<f64> {
|
||
let n = self.ndim * self.ndim;
|
||
let mut matrix = Array2::zeros([self.ndim + 1, self.ndim + 1]);
|
||
for (m, p) in matrix
|
||
.slice_mut(s![..self.ndim, ..self.ndim])
|
||
.iter_mut()
|
||
.zip_eq(self.dparameters.iter().take(n))
|
||
{
|
||
*m = *p;
|
||
}
|
||
for (m, p) in matrix
|
||
.slice_mut(s![..self.ndim, self.ndim])
|
||
.iter_mut()
|
||
.zip_eq(self.dparameters.iter().skip(n))
|
||
{
|
||
*m = *p;
|
||
}
|
||
matrix
|
||
}
|
||
|
||
/// get the inverse transform
|
||
pub fn inverse(&self) -> Result<Self, Error> {
|
||
let matrix = self.matrix();
|
||
let inverse = matrix.inv()?;
|
||
let parameters = inverse
|
||
.slice(s![..self.ndim, ..self.ndim])
|
||
.iter()
|
||
.chain(inverse.slice(s![..self.ndim, self.ndim]))
|
||
.cloned()
|
||
.collect();
|
||
|
||
Ok(Self {
|
||
parameters,
|
||
dparameters: vec![0f64; self.ndim * self.ndim + self.ndim],
|
||
center: self.center.clone(),
|
||
shape: self.shape.clone(),
|
||
ndim: self.ndim,
|
||
dimension: PhantomData,
|
||
})
|
||
}
|
||
|
||
/// adapt the transform to a new center and shape
|
||
pub fn adapt(&mut self, center: &[f64], shape: &[usize]) {
|
||
self.center = self
|
||
.shape
|
||
.iter()
|
||
.zip_eq(shape.iter())
|
||
.zip_eq(center)
|
||
.map(|((a, b), o)| o + (((a - b) as f64) / 2.0))
|
||
.collect();
|
||
self.shape = shape.to_vec();
|
||
}
|
||
|
||
pub fn interpolate<'a, const B: usize, T, A>(&self, image: A) -> Result<Array<f64, D>, Error>
|
||
where
|
||
A: AsArray<'a, T, D>,
|
||
D: Dimension,
|
||
T: 'a + Clone + AsPrimitive<f64>,
|
||
BSpline<B, D>: BSplineTrait<D>,
|
||
{
|
||
let image = image.into();
|
||
let bspline = BSpline::<B, _>::new(&image);
|
||
bspline.interpolate(self)
|
||
}
|
||
|
||
pub fn interpolate_par<'a, const B: usize, T, A>(
|
||
&self,
|
||
image: A,
|
||
) -> Result<Array<f64, D>, Error>
|
||
where
|
||
A: AsArray<'a, T, D>,
|
||
D: Dimension,
|
||
T: 'a + Clone + AsPrimitive<f64>,
|
||
BSpline<B, D>: BSplineTrait<D>,
|
||
{
|
||
let image = image.into();
|
||
let bspline = BSpline::<B, _>::new(&image);
|
||
bspline.interpolate_par(self)
|
||
}
|
||
}
|
||
|
||
impl Transform<Ix2> {
|
||
/// create a transform from a rotation in radians
|
||
pub fn from_rotation(theta: f64, center: &[f64]) -> Self {
|
||
Self {
|
||
parameters: vec![
|
||
theta.cos(),
|
||
-theta.sin(),
|
||
theta.sin(),
|
||
theta.cos(),
|
||
0.0,
|
||
0.0,
|
||
],
|
||
dparameters: vec![0f64; 6],
|
||
center: center.to_vec(),
|
||
shape: vec![0usize; 2],
|
||
ndim: 2,
|
||
dimension: PhantomData,
|
||
}
|
||
}
|
||
|
||
/// multiply self with a rotation matrix
|
||
pub fn with_rotation(&self, theta: f64) -> Self {
|
||
Self::from_rotation(theta, &self.center) * self
|
||
}
|
||
|
||
/// multiply self with a rotation matrix
|
||
pub fn with_rotation_around(&self, theta: f64, center: &[f64]) -> Self {
|
||
Self::from_rotation(theta, center) * self
|
||
}
|
||
}
|
||
|
||
#[cfg(test)]
|
||
mod tests {
|
||
use crate::julia_image;
|
||
use crate::transform::Transform;
|
||
use itertools::Itertools;
|
||
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<dyn std::error::Error>> {
|
||
let image_a = julia_image(
|
||
&[8000, 6000],
|
||
&[1.0, 0.0, 0.0, 1.0, 0.0, 0.0],
|
||
&[3999.5, 2999.5],
|
||
&[-0.8, 0.156],
|
||
);
|
||
let transform = Transform::from_rotation(1.0, &[3999.5, 2999.5]);
|
||
let image_b = transform.interpolate_par::<3, _, _>(image_a.view())?;
|
||
let image_c = transform.interpolate::<3, _, _>(image_a.view())?;
|
||
|
||
assert!(
|
||
image_b
|
||
.into_iter()
|
||
.zip_eq(image_c.into_iter())
|
||
.all(|(a, b)| a == b)
|
||
);
|
||
Ok(())
|
||
}
|
||
|
||
#[test]
|
||
fn interpolate_serial() -> Result<(), Box<dyn std::error::Error>> {
|
||
let image_a = julia_image(
|
||
&[8000, 6000],
|
||
&[1.0, 0.0, 0.0, 1.0, 0.0, 0.0],
|
||
&[3999.5, 2999.5],
|
||
&[-0.8, 0.156],
|
||
);
|
||
let transform = Transform::from_rotation(1.0, &[3999.5, 2999.5]);
|
||
let image_b = transform.interpolate::<3, _, _>(image_a.view())?;
|
||
assert_ne!(image_b[[4000, 3000]], 0.0);
|
||
Ok(())
|
||
}
|
||
|
||
#[test]
|
||
fn interpolate_par() -> Result<(), Box<dyn std::error::Error>> {
|
||
let image_a = julia_image(
|
||
&[8000, 6000],
|
||
&[1.0, 0.0, 0.0, 1.0, 0.0, 0.0],
|
||
&[3999.5, 2999.5],
|
||
&[-0.8, 0.156],
|
||
);
|
||
let transform = Transform::from_rotation(1.0, &[3999.5, 2999.5]);
|
||
let image_b = transform.interpolate_par::<3, _, _>(image_a.view())?;
|
||
assert_ne!(image_b[[4000, 3000]], 0.0);
|
||
Ok(())
|
||
}
|
||
|
||
#[test]
|
||
fn grad_check() -> Result<(), Box<dyn std::error::Error>> {
|
||
use crate::bspline::{BSpline, BSplineTrait};
|
||
use crate::metric::{MattesMetric, SamplingArg};
|
||
use algos::ObjectiveFunction;
|
||
let im_a = julia_image(
|
||
&[100, 1],
|
||
&[1.0, 0.0, 0.0, 0.01, 0.0, 0.0],
|
||
&[99.5, 0.5],
|
||
&[-0.8, 0.156],
|
||
)
|
||
.slice(s![.., 0])
|
||
.mapv(|i| i as f64);
|
||
let q = vec![0.85, 4.0];
|
||
let im_b =
|
||
Transform::new(q.clone(), vec![im_a.shape()[0]]).interpolate::<1, _, _>(&im_a)?;
|
||
|
||
// Test with ALL points (like grad_check uses)
|
||
let metric_all = MattesMetric::<ndarray::Ix1>::new(
|
||
BSpline::new(im_b.view()),
|
||
BSpline::new(im_a.view()),
|
||
SamplingArg::FixedAt((0..100).map(|i| vec![i as f64]).collect()),
|
||
3,
|
||
0.05,
|
||
)?
|
||
.with_fixed_mu(crate::metric::FixedMu::new_none(1));
|
||
|
||
let identity = vec![1.0, 0.0];
|
||
let val_i = metric_all.evaluate(&identity);
|
||
let grad_i = metric_all.gradient(&identity).unwrap();
|
||
println!("ALL POINTS - identity: val={}, grad={:?}", val_i, grad_i);
|
||
let truth = vec![0.85, 4.0];
|
||
let val_t = metric_all.evaluate(&truth);
|
||
let grad_t = metric_all.gradient(&truth).unwrap();
|
||
println!("ALL POINTS - truth: val={}, grad={:?}", val_t, grad_t);
|
||
let inv = vec![1.0 / 0.85, -4.0 / 0.85];
|
||
let val_inv = metric_all.evaluate(&inv);
|
||
println!("ALL POINTS - inverse: val={}", val_inv);
|
||
|
||
// Test with 100 random points (like register1 uses at finest level)
|
||
let metric_rand = MattesMetric::<ndarray::Ix1>::new(
|
||
BSpline::new(im_b.view()),
|
||
BSpline::new(im_a.view()),
|
||
SamplingArg::Fixed(100),
|
||
3,
|
||
0.05,
|
||
)?
|
||
.with_fixed_mu(crate::metric::FixedMu::new_none(1));
|
||
|
||
let val_i2 = metric_rand.evaluate(&identity);
|
||
let grad_i2 = metric_rand.gradient(&identity).unwrap();
|
||
println!("RAND 100 - identity: val={}, grad={:?}", val_i2, grad_i2);
|
||
|
||
// Test with 100 random points and 32 bins (like level 0)
|
||
let metric_32 = MattesMetric::<ndarray::Ix1>::new(
|
||
BSpline::new(im_b.view()),
|
||
BSpline::new(im_a.view()),
|
||
SamplingArg::Fixed(100),
|
||
32,
|
||
0.05,
|
||
)?
|
||
.with_fixed_mu(crate::metric::FixedMu::new_none(1));
|
||
|
||
let val_i3 = metric_32.evaluate(&identity);
|
||
let grad_i3 = metric_32.gradient(&identity).unwrap();
|
||
println!(
|
||
"RAND 100 bins=32 - identity: val={}, grad={:?}",
|
||
val_i3, grad_i3
|
||
);
|
||
|
||
let eps = 1e-5;
|
||
for i in 0..2 {
|
||
let mut p_plus = identity.clone();
|
||
let mut p_minus = identity.clone();
|
||
p_plus[i] += eps;
|
||
p_minus[i] -= eps;
|
||
let num_grad =
|
||
(metric_all.evaluate(&p_plus) - metric_all.evaluate(&p_minus)) / (2.0 * eps);
|
||
println!(
|
||
"numerical d/dmu[{}] = {} (analytical: {})",
|
||
i, num_grad, grad_i[i]
|
||
);
|
||
}
|
||
Ok(())
|
||
}
|
||
|
||
#[test]
|
||
fn metric_landscape() -> Result<(), Box<dyn std::error::Error>> {
|
||
use crate::bspline::{BSpline, BSplineTrait};
|
||
use crate::metric::{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],
|
||
&[99.5, 0.5],
|
||
&[-0.8, 0.156],
|
||
)
|
||
.slice(s![.., 0])
|
||
.mapv(|i| i as f64);
|
||
let im_b =
|
||
Transform::new(vec![0.85, 4.0], vec![im_a.shape()[0]]).interpolate::<1, _, _>(&im_a)?;
|
||
|
||
// Smooth both images with sigma=8 like level 0
|
||
let sigma = Sigma::Absolute(vec![8.0]);
|
||
let sf = sigma.smooth(im_a.view())?;
|
||
let sm = sigma.smooth(im_b.view())?;
|
||
|
||
let points: Vec<Vec<f64>> = (0..100).map(|i| vec![i as f64]).collect();
|
||
|
||
let metric = MattesMetric::<ndarray::Ix1>::new(
|
||
BSpline::new(sf.view()),
|
||
BSpline::new(sm.view()),
|
||
SamplingArg::FixedAt(points.clone()),
|
||
32,
|
||
0.05,
|
||
)?
|
||
.with_fixed_mu(crate::metric::FixedMu::new_none(1));
|
||
|
||
// 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]),
|
||
];
|
||
|
||
println!("=== SMOOTHED ALL 100 integer points bins=32 ===");
|
||
for (name, p) in &test_points {
|
||
let val = metric.evaluate(p);
|
||
let grad = metric.gradient(p);
|
||
match grad {
|
||
Some(g) => println!(" {}: val={:.6}, grad={:?}", name, val, g),
|
||
None => println!(" {}: val={:.6}, grad=None", name, val),
|
||
}
|
||
}
|
||
|
||
// Now test with Fixed(100) random points
|
||
let metric_rand = MattesMetric::<ndarray::Ix1>::new(
|
||
BSpline::new(sf.view()),
|
||
BSpline::new(sm.view()),
|
||
SamplingArg::Fixed(100),
|
||
32,
|
||
0.05,
|
||
)?
|
||
.with_fixed_mu(crate::metric::FixedMu::new_none(1));
|
||
|
||
println!("=== SMOOTHED Random 100 points bins=32 ===");
|
||
for (name, p) in &test_points {
|
||
let val = metric_rand.evaluate(p);
|
||
let grad = metric_rand.gradient(p);
|
||
match grad {
|
||
Some(g) => println!(" {}: val={:.6}, grad={:?}", name, val, g),
|
||
None => println!(" {}: val={:.6}, grad=None", name, val),
|
||
}
|
||
}
|
||
|
||
Ok(())
|
||
}
|
||
|
||
#[test]
|
||
fn register1() -> Result<(), Box<dyn std::error::Error>> {
|
||
let im_a = julia_image(
|
||
&[100, 1],
|
||
&[1.0, 0.0, 0.0, 0.01, 0.0, 0.0],
|
||
&[99.5, 0.5],
|
||
&[-0.8, 0.156],
|
||
)
|
||
.slice(s![.., 0])
|
||
.mapv(|i| i as f64);
|
||
let q = vec![0.85, 2.0];
|
||
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 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
|
||
);
|
||
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 shape = [200, 200];
|
||
let center = [99.5, 99.5];
|
||
|
||
let fixed = julia_image(
|
||
&shape,
|
||
&[1.0, 0.0, 0.0, 1.0, 0.0, 0.0],
|
||
¢er,
|
||
&[-0.8, 0.156],
|
||
)
|
||
.mapv(|i| i as f64);
|
||
|
||
// random rotation ±25°
|
||
let angle: f64 = rng.random_range(-25.0f64..25.0).to_radians();
|
||
let (s, c) = angle.sin_cos();
|
||
|
||
// random scale 0.85–1.15
|
||
let sx = rng.random_range(0.85..1.15);
|
||
let sy = rng.random_range(0.85..1.15);
|
||
|
||
// small shear
|
||
let shx: f64 = rng.random_range(-0.1..0.1);
|
||
let shy: f64 = rng.random_range(-0.1..0.1);
|
||
|
||
// translation ±30 px
|
||
let tx: f64 = rng.random_range(-30.0..30.0);
|
||
let ty: f64 = rng.random_range(-30.0..30.0);
|
||
|
||
// [m00, m01, m10, m11, tx, ty]
|
||
let params = vec![c * sx, -s * sy + shx, s * sx + shy, c * sy, tx, ty];
|
||
|
||
let transform =
|
||
Transform::<Ix2>::new_with_center(params.clone(), center.to_vec(), shape.to_vec());
|
||
let moving = transform.interpolate::<3, _, _>(fixed.view())?;
|
||
|
||
let q_inv = transform.inverse()?.parameters;
|
||
|
||
let (t, _steps) = Transform::<Ix2>::register_debug(
|
||
fixed.view(),
|
||
moving.view(),
|
||
vec![None, None, None, None, None, None],
|
||
None,
|
||
None,
|
||
)?;
|
||
println!("params: {:?}", params);
|
||
println!("result: {:?}", t.parameters);
|
||
println!("q_inv: {:?}", q_inv);
|
||
let sse: f64 = t
|
||
.parameters
|
||
.iter()
|
||
.zip(q_inv.iter())
|
||
.map(|(a, b)| (a - b).powi(2))
|
||
.sum();
|
||
println!("sse: {sse}");
|
||
assert!(sse < 1.0);
|
||
Ok(())
|
||
}
|
||
|
||
#[test]
|
||
fn metric_landscape_2d() -> Result<(), Box<dyn std::error::Error>> {
|
||
use crate::bspline::{BSpline, BSplineTrait};
|
||
use crate::filter::gaussian_smooth;
|
||
use crate::metric::{MattesMetric, SamplingArg};
|
||
use algos::ObjectiveFunction;
|
||
|
||
let shape = [200, 200];
|
||
let center = [99.5, 99.5];
|
||
let im_a = julia_image(
|
||
&shape,
|
||
&[1.0, 0.0, 0.0, 1.0, 0.0, 0.0],
|
||
¢er,
|
||
&[-0.8, 0.156],
|
||
)
|
||
.mapv(|i| i as f64);
|
||
let rotation = Transform::<Ix2>::from_rotation(f64::PI() / 4.0, ¢er);
|
||
let im_b = rotation.interpolate::<3, _, _>(im_a.view())?;
|
||
|
||
let q_inv = rotation.inverse()?.parameters;
|
||
let p = rotation.parameters.clone();
|
||
let identity = vec![1.0, 0.0, 0.0, 1.0, 0.0, 0.0];
|
||
let edge = 0.01;
|
||
|
||
// Test at full resolution with different sigma levels (matching elastix approach)
|
||
// Elastix does NOT downsample — it smooths at full resolution with
|
||
// sigma = 0.5 * factor where factor is from the pyramid schedule [8, 4, 2, 1]
|
||
println!(
|
||
"=== Full resolution (no downsampling, matching elastix FixedSmoothingImagePyramid) ==="
|
||
);
|
||
for (level, sigma_val) in [4.0, 2.0, 1.0, 0.5].iter().enumerate() {
|
||
let f = gaussian_smooth(im_a.view(), &[*sigma_val; 2])?;
|
||
let m = gaussian_smooth(im_b.view(), &[*sigma_val; 2])?;
|
||
let bf = BSpline::<0, _>::new(f.view());
|
||
let bm = BSpline::<3, _>::new(m.view());
|
||
let metric = MattesMetric::new(bf, bm, SamplingArg::Random(3000), 128, edge)?;
|
||
|
||
let mi_id = metric.evaluate(&identity);
|
||
let mi_p = metric.evaluate(&p);
|
||
let mi_qinv = metric.evaluate(&q_inv);
|
||
|
||
println!(
|
||
" Level {} (sigma={:.1}, {}x{}): id={:.4} p={:.4} q_inv={:.4} → {}",
|
||
level,
|
||
sigma_val,
|
||
f.shape()[0],
|
||
f.shape()[1],
|
||
-mi_id,
|
||
-mi_p,
|
||
-mi_qinv,
|
||
if -mi_qinv > -mi_p {
|
||
"Q_INV correct"
|
||
} else {
|
||
"P incorrect (landscape inverted!)"
|
||
}
|
||
);
|
||
}
|
||
|
||
// Verify: at every sigma level, q_inv should have higher MI than p
|
||
let f4 = gaussian_smooth(im_a.view(), &[4.0, 4.0])?;
|
||
let m4 = gaussian_smooth(im_b.view(), &[4.0, 4.0])?;
|
||
let bf4 = BSpline::<0, _>::new(f4.view());
|
||
let bm4 = BSpline::<3, _>::new(m4.view());
|
||
let metric4 = MattesMetric::new(bf4, bm4, SamplingArg::Random(5000), 128, edge)?;
|
||
let mi_p_coarse = metric4.evaluate(&p);
|
||
let mi_qinv_coarse = metric4.evaluate(&q_inv);
|
||
println!(
|
||
"\nCoarsest (sigma=4.0): p={:.6} q_inv={:.6}",
|
||
-mi_p_coarse, -mi_qinv_coarse
|
||
);
|
||
assert!(
|
||
-mi_qinv_coarse > -mi_p_coarse,
|
||
"MI landscape is inverted at coarsest level! q_inv={} should be > p={}",
|
||
-mi_qinv_coarse,
|
||
-mi_p_coarse
|
||
);
|
||
|
||
Ok(())
|
||
}
|
||
|
||
#[test]
|
||
fn angle_sweep() -> Result<(), Box<dyn std::error::Error>> {
|
||
use crate::bspline::{BSpline, BSplineTrait};
|
||
use crate::filter::gaussian_smooth;
|
||
use crate::metric::{MattesMetric, SamplingArg};
|
||
use algos::ObjectiveFunction;
|
||
let shape = [200, 200];
|
||
let center = [99.5, 99.5];
|
||
let im_a = julia_image(
|
||
&shape,
|
||
&[1.0, 0.0, 0.0, 1.0, 0.0, 0.0],
|
||
¢er,
|
||
&[-0.8, 0.156],
|
||
)
|
||
.mapv(|i| i as f64);
|
||
let edge = 0.01;
|
||
// Full resolution, sigma=4.0 — matching elastix FixedSmoothingImagePyramid (no downsampling)
|
||
for angle_deg in [5.0f64, 10.0, 15.0, 20.0, 25.0, 30.0, 35.0, 40.0, 45.0] {
|
||
let rotation = Transform::<Ix2>::from_rotation(angle_deg.to_radians(), ¢er);
|
||
let im_b = rotation.interpolate::<3, _, _>(im_a.view())?;
|
||
let q_inv = rotation.inverse()?.parameters;
|
||
let p = rotation.parameters.clone();
|
||
let f = gaussian_smooth(im_a.view(), &[4.0, 4.0])?;
|
||
let m = gaussian_smooth(im_b.view(), &[4.0, 4.0])?;
|
||
let metric = MattesMetric::new(
|
||
BSpline::<0, _>::new(f.view()),
|
||
BSpline::<3, _>::new(m.view()),
|
||
SamplingArg::Random(3000),
|
||
128,
|
||
edge,
|
||
)?;
|
||
let id = vec![1.0, 0.0, 0.0, 1.0, 0.0, 0.0];
|
||
let mi_id = metric.evaluate(&id);
|
||
let mi_p = metric.evaluate(&p);
|
||
let mi_qi = metric.evaluate(&q_inv);
|
||
println!(
|
||
"{:5.1}°: id={:.3} p={:.3} q_inv={:.3} → {}",
|
||
angle_deg,
|
||
-mi_id,
|
||
-mi_p,
|
||
-mi_qi,
|
||
if mi_p > mi_qi { "P wins" } else { "Q_INV wins" }
|
||
);
|
||
}
|
||
Ok(())
|
||
}
|
||
|
||
fn read_tiff<P: AsRef<Path>>(path: P) -> Result<Array2<f64>, Box<dyn std::error::Error>> {
|
||
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::<Vec<_>>(),
|
||
DecodingResult::U16(data) => data.into_iter().map(|i| i as f64).collect::<Vec<_>>(),
|
||
DecodingResult::U32(data) => data.into_iter().map(|i| i as f64).collect::<Vec<_>>(),
|
||
DecodingResult::U64(data) => data.into_iter().map(|i| i as f64).collect::<Vec<_>>(),
|
||
DecodingResult::I8(data) => data.into_iter().map(|i| i as f64).collect::<Vec<_>>(),
|
||
DecodingResult::I16(data) => data.into_iter().map(|i| i as f64).collect::<Vec<_>>(),
|
||
DecodingResult::I32(data) => data.into_iter().map(|i| i as f64).collect::<Vec<_>>(),
|
||
DecodingResult::I64(data) => data.into_iter().map(|i| i as f64).collect::<Vec<_>>(),
|
||
DecodingResult::F16(data) => data.into_iter().map(f64::from).collect::<Vec<_>>(),
|
||
DecodingResult::F32(data) => data.into_iter().map(|i| i as f64).collect::<Vec<_>>(),
|
||
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<dyn std::error::Error>> {
|
||
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::<Ix2>::register(
|
||
fixed.view(),
|
||
moving.view(),
|
||
vec![None, None, None, None, None, None],
|
||
None,
|
||
None,
|
||
)?;
|
||
|
||
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()
|
||
.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)?;
|
||
|
||
assert!(mean_diff < 0.1);
|
||
|
||
Ok(())
|
||
}
|
||
|
||
#[test]
|
||
fn register_real_images2() -> Result<(), Box<dyn std::error::Error>> {
|
||
let fixed = read_tiff("test_files/fixed.tif")?;
|
||
let e = Transform::<Ix2>::new(vec![0.8, 0.0, 0.0, 1.0, 0.0, 0.0], fixed.shape().to_vec())
|
||
.inverse()?;
|
||
let moving = e.interpolate::<3, _, _>(fixed.view())?;
|
||
|
||
let t = Transform::<Ix2>::register(
|
||
fixed.view(),
|
||
moving.view(),
|
||
vec![None, None, None, None, None, None],
|
||
None,
|
||
None,
|
||
)?;
|
||
|
||
let e_inv = e.inverse()?;
|
||
|
||
// 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(moving.mapv(|i| i as u16), 2, 0, 0)?;
|
||
|
||
assert!(mean_diff < 0.1);
|
||
|
||
Ok(())
|
||
}
|
||
}
|