Files
sitk-registration-sys/src/registration.rs
T

646 lines
21 KiB
Rust

//! Some structs and methods to make working with registration and interpolation methods in
//! SimpleITK more Rust friendly.
use crate::simple;
use anyhow::{Result, anyhow};
use autocxx::prelude::*;
use cxx::{CxxVector, UniquePtr, let_cxx_string};
use ndarray::{Array2, ArrayView2, AsArray, Ix2, array, s};
use serde::{Deserialize, Serialize};
use serde_yaml::{from_reader, to_writer};
use std::fs::File;
use std::marker::PhantomData;
use std::ops::{Deref, DerefMut, Mul};
use std::path::PathBuf;
use num::Complex;
use tempfile::tempdir;
/// a trait marking number types that can be used in sitk:
/// (u/i)(8/16/32/64), (u/i)size, f(32/64)
pub trait PixelType: Clone {
const PT: simple::PixelIDValueEnum;
}
macro_rules! pixel_type_impl {
($($T:ty: $sitk:expr $(,)?)*) => {
$(
impl PixelType for $T {
const PT: simple::PixelIDValueEnum = $sitk;
}
)*
};
}
pixel_type_impl! {
u8: simple::PixelIDValueEnum::sitkUInt8,
i8: simple::PixelIDValueEnum::sitkInt8,
u16: simple::PixelIDValueEnum::sitkUInt16,
i16: simple::PixelIDValueEnum::sitkInt16,
u32: simple::PixelIDValueEnum::sitkUInt32,
i32: simple::PixelIDValueEnum::sitkInt32,
u64: simple::PixelIDValueEnum::sitkUInt64,
i64: simple::PixelIDValueEnum::sitkInt64,
f32: simple::PixelIDValueEnum::sitkFloat32,
f64: simple::PixelIDValueEnum::sitkFloat64,
}
#[cfg(target_pointer_width = "64")]
pixel_type_impl!(usize: simple::PixelIDValueEnum::sitkUInt64);
#[cfg(target_pointer_width = "32")]
pixel_type_impl!(usize: simple::PixelIDValueEnum::sitkUInt32);
#[cfg(target_pointer_width = "64")]
pixel_type_impl!(isize: simple::PixelIDValueEnum::sitkInt64);
#[cfg(target_pointer_width = "32")]
pixel_type_impl!(isize: simple::PixelIDValueEnum::sitkInt32);
/// Struct holding a pointer to an image
pub struct Image<T: PixelType> {
image: UniquePtr<simple::Image>,
pixel_type: PhantomData<T>,
}
impl<T: PixelType> Deref for Image<T> {
type Target = UniquePtr<simple::Image>;
fn deref(&self) -> &Self::Target {
&self.image
}
}
impl<T: PixelType> Image<T> {
/// encapsulate an itk::simple::Image
pub fn new(image: UniquePtr<simple::Image>) -> Self {
Self {
image,
pixel_type: PhantomData,
}
}
/// take an ndarray Array2 and turn it into a SimpleITK image
pub fn from_array<'a, A>(array: A) -> Self
where
T: 'a + PixelType,
A: AsArray<'a, T, Ix2>,
{
let array = array.into();
let shape = array.shape();
let width = (shape[1] as u32).into();
let height = (shape[0] as u32).into();
let mut image = simple::Image::new3(width, height, T::PT).within_unique_ptr();
image.pin_mut().MakeUnique();
let buffer = image.pin_mut().GetBufferAsVoid();
unsafe { std::ptr::copy(array.as_ptr(), buffer as *mut T, shape[0] * shape[1]) };
Self {
image,
pixel_type: PhantomData,
}
}
/// return as an ndarray Array2
pub fn as_array(&self) -> Array2<T> {
let width = u32::from(self.image.GetWidth()) as usize;
let height = u32::from(self.image.GetHeight()) as usize;
let mut array = Array2::<T>::uninit((height, width));
let buffer = self.image.GetBufferAsVoid1();
unsafe {
std::ptr::copy(
buffer as *const T,
array.as_mut_ptr() as *mut T,
width * height,
);
array.assume_init()
}
}
}
impl<'a, A, T> From<A> for Image<T>
where
T: 'a + PixelType,
A: AsArray<'a, T, Ix2>,
{
fn from(value: A) -> Self {
Self::from_array(value.into())
}
}
impl<T> From<Image<T>> for Array2<T>
where
T: PixelType,
{
fn from(value: Image<T>) -> Self {
value.as_array()
}
}
/// a struct describing the transform
#[derive(Clone, Debug, Deserialize, Serialize, PartialEq)]
pub struct AffineTransform {
/// flattened 2x2 rotation matrix + translation
pub parameters: [f64; 6],
/// error / significance on parameters
pub dparameters: [f64; 6],
/// the point about which rotations are performed
pub origin: [f64; 2],
/// the shape of images for which this transform is meant
pub shape: [usize; 2],
}
impl Mul for AffineTransform {
type Output = AffineTransform;
#[allow(clippy::suspicious_arithmetic_impl)]
fn mul(self, other: AffineTransform) -> AffineTransform {
let m = self.matrix().dot(&other.matrix());
let dm = self.dmatrix().dot(&other.matrix()) + self.matrix().dot(&other.dmatrix());
AffineTransform {
parameters: [
m[[0, 0]],
m[[0, 1]],
m[[1, 0]],
m[[1, 1]],
m[[2, 0]],
m[[2, 1]],
],
dparameters: [
dm[[0, 0]],
dm[[0, 1]],
dm[[1, 0]],
dm[[1, 1]],
dm[[2, 0]],
dm[[2, 1]],
],
origin: self.origin,
shape: self.shape,
}
}
}
impl Eq for AffineTransform {}
impl AffineTransform {
/// parameters: flat 2x2 part of matrix, translation; origin: center of rotation
pub fn new(parameters: [f64; 6], origin: [f64; 2], shape: [usize; 2]) -> Self {
Self {
parameters,
dparameters: [0f64; 6],
origin,
shape,
}
}
/// find the affine transform which transforms moving into fixed
pub fn register_affine<F, M, T>(fixed: F, moving: M) -> Result<AffineTransform>
where
F: Into<Image<T>>,
M: Into<Image<T>>,
T: PixelType,
{
Self::register(fixed, moving, true)
}
/// find the translation which transforms moving into fixed
pub fn register_translation<F, M, T>(fixed: F, moving: M) -> Result<AffineTransform>
where
F: Into<Image<T>>,
M: Into<Image<T>>,
T: PixelType,
{
Self::register(fixed, moving, false)
}
/// find the transform which transforms moving into fixed
pub fn register<F, M, T>(fixed: F, moving: M, affine: bool) -> Result<AffineTransform>
where
F: Into<Image<T>>,
M: Into<Image<T>>,
T: PixelType,
{
let tmp_folder = tempdir()?;
let fixed = fixed.into();
let moving = moving.into();
let width = u32::from(fixed.GetWidth()) as usize;
let height = u32::from(fixed.GetHeight()) as usize;
let_cxx_string!(transform_name = if affine { "affine" } else { "translation" });
let parameter_map = crate::ffi_extra::get_default_parameter_map(&transform_name);
let mut tfilter = simple::ElastixImageFilter::new().within_box();
tfilter.as_mut().LogToConsoleOff();
tfilter.as_mut().LogToFileOff();
tfilter.as_mut().SetLogToFile(false);
tfilter.as_mut().SetFixedImage(&fixed);
tfilter.as_mut().SetMovingImage(&moving);
crate::ffi_extra::set_parameter_map(&mut tfilter, &parameter_map);
tfilter.as_mut().SetParameter("WriteResultImage", "False");
tfilter
.as_mut()
.SetOutputDirectory(tmp_folder.path().display().to_string());
let _ = tfilter.as_mut().Execute().within_unique_ptr();
let_cxx_string!(tp = "TransformParameters");
let p = crate::ffi_extra::get_transform_parameter_map(tfilter.as_mut().deref_mut(), 0)
.get(&tp)
.iter()
.map(|i| i.to_string_lossy().parse::<f64>())
.collect::<Result<Vec<f64>, _>>()?;
let parameters = if affine {
[p[0], p[1], p[2], p[3], p[4], p[5]]
} else {
[1.0, 0.0, 0.0, 1.0, p[0], p[1]]
};
let origin = [((height - 1) as f64) / 2.0, ((width - 1) as f64) / 2.0];
let shape = [height, width];
Ok(AffineTransform::new(parameters, origin, shape))
}
/// create a transform from a xy translation
pub fn from_translation(translation: [f64; 2]) -> Self {
AffineTransform {
parameters: [1f64, 0f64, 0f64, 1f64, translation[0], translation[1]],
dparameters: [0f64; 6],
origin: [0f64; 2],
shape: [0usize; 2],
}
}
/// read a transform from a file
pub fn from_file(path: PathBuf) -> Result<Self> {
let file = File::open(path)?;
Ok(from_reader(file)?)
}
/// write a transform to a file
pub fn to_file(&self, path: PathBuf) -> Result<()> {
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 {
self.parameters == [1f64, 0f64, 0f64, 1f64, 0f64, 0f64]
}
/// transform an image using nearest neighbor interpolation
pub fn transform_image_bspline<I, T>(&self, image: I) -> Result<Image<T>>
where
I: Into<Image<T>>,
T: PixelType,
{
Self::interp(self, image, false)
}
/// transform an image using bspline interpolation
pub fn transform_image_nearest_neighbor<I, T>(&self, image: I) -> Result<Image<T>>
where
I: Into<Image<T>>,
T: PixelType,
{
Self::interp(self, image, true)
}
/// transform an image
pub fn interp<I, T>(&self, image: I, nearest_neighbor: bool) -> Result<Image<T>>
where
I: Into<Image<T>>,
T: PixelType,
{
let image = image.into();
let width = u32::from(image.GetWidth()) as usize;
let height = u32::from(image.GetHeight()) as usize;
let origin = [((width - 1) as f64) / 2f64, ((height - 1) as f64) / 2f64];
let p = self.parameters;
let matrix = cxx_vector([p[0], p[1], p[2], p[3]]);
let translation = cxx_vector([p[4], p[5]]);
let fixed_center = cxx_vector(origin);
let affine_transform =
simple::AffineTransform::new3(&matrix, &translation, &fixed_center).within_unique_ptr();
let transform = <_ as AsRef<simple::Transform>>::as_ref(affine_transform.as_ref().unwrap());
let interpolator = if nearest_neighbor {
simple::InterpolatorEnum::sitkBSpline
} else {
simple::InterpolatorEnum::sitkNearestNeighbor
};
Ok(Image::<T>::new(
simple::Resample(
&image,
transform,
interpolator,
0.0,
simple::PixelIDValueEnum::sitkUnknown,
nearest_neighbor,
)
.within_unique_ptr(),
))
}
/// get coordinates resulting from transforming input coordinates, coordinates must have two
/// columns: x & y
pub fn transform_coordinates<'a, A, T>(&self, coordinates: A) -> Result<Array2<f64>>
where
T: 'a + Clone + Into<f64>,
A: AsArray<'a, T, Ix2>,
{
let coordinates = coordinates.into();
let s = coordinates.shape();
if s[1] != 2 {
return Err(anyhow!("coordinates must have two columns"));
}
let m = self.matrix();
let mut res = Array2::zeros([s[0], s[1]]);
for i in 0..s[0] {
let a = array![
coordinates[[i, 0]].clone().into(),
coordinates[[i, 1]].clone().into(),
1f64
]
.to_owned();
let b = m.dot(&a);
res.slice_mut(s![i, ..]).assign(&b.slice(s![..2]));
}
Ok(res)
}
/// get the matrix defining the transform
pub fn matrix(&self) -> Array2<f64> {
Array2::from_shape_vec(
(3, 3),
vec![
self.parameters[0],
self.parameters[1],
self.parameters[4],
self.parameters[2],
self.parameters[3],
self.parameters[5],
0f64,
0f64,
1f64,
],
)
.unwrap()
}
/// get the matrix describing the error of the transform
pub fn dmatrix(&self) -> Array2<f64> {
Array2::from_shape_vec(
(3, 3),
vec![
self.dparameters[0],
self.dparameters[1],
self.dparameters[4],
self.dparameters[2],
self.dparameters[3],
self.dparameters[5],
0f64,
0f64,
1f64,
],
)
.unwrap()
}
/// get the inverse transform
pub fn inverse(&self) -> Result<AffineTransform> {
fn det(a: ArrayView2<f64>) -> f64 {
(a[[0, 0]] * a[[1, 1]]) - (a[[0, 1]] * a[[1, 0]])
}
let m = self.matrix();
let d = det(m.slice(s![..2, ..2]));
if d == 0f64 {
return Err(anyhow!("transform matrix is not invertible"));
}
let parameters = [
det(m.slice(s![1.., 1..])) / d,
-det(m.slice(s![..;2, 1..])) / d,
-det(m.slice(s![1.., ..;2])) / d,
det(m.slice(s![..;2, ..;2])) / d,
det(m.slice(s![..2, 1..])) / d,
-det(m.slice(s![..2, ..;2])) / d,
];
Ok(AffineTransform {
parameters,
dparameters: [0f64; 6],
origin: self.origin,
shape: self.shape,
})
}
/// adapt the transform to a new origin and shape
pub fn adapt(&mut self, origin: [f64; 2], shape: [usize; 2]) {
self.origin = [
origin[0] + (((self.shape[0] - shape[0]) as f64) / 2f64),
origin[1] + (((self.shape[1] - shape[1]) as f64) / 2f64),
];
self.shape = shape;
}
}
/// conveniently collect an iterator into a CxxVector
pub fn cxx_vector<T, I>(vec: I) -> UniquePtr<CxxVector<T>>
where
I: IntoIterator<Item = T>,
T: cxx::vector::VectorElement + cxx::ExternType<Kind = cxx::kind::Trivial>,
{
let mut v = CxxVector::new();
v.pin_mut().extend(vec);
v
}
/// An example of generating julia fractals, for testing purposes.
pub fn julia_image(shift_x: f32, shift_y: f32) -> Result<Array2<u8>> {
let imgx = 800;
let imgy = 600;
let scalex = 3.0 / imgx as f32;
let scaley = 3.0 / imgy as f32;
let mut im = Array2::<u8>::zeros((imgy, imgx));
for x in 0..imgx {
for y in 0..imgy {
let cy = (y as f32 + shift_y) * scalex - 1.5;
let cx = (x as f32 + shift_x) * scaley - 1.5;
let c = Complex::new(-0.4, 0.6);
let mut z = Complex::new(cy, cx);
let mut i = 0;
while i < 255 && z.norm() <= 2.0 {
z = z * z + c;
i += 1;
}
im[[y, x]] = i as u8;
}
}
Ok(im)
}
#[cfg(test)]
mod tests {
use super::*;
use anyhow::Result;
use ndarray::Array2;
use tempfile::NamedTempFile;
#[test]
fn test_serialization() -> Result<()> {
let file = NamedTempFile::new()?;
let t = AffineTransform::new([1.2, 0.3, -0.4, 0.9, 10.2, -9.5], [59.5, 49.5], [120, 100]);
t.to_file(file.path().to_path_buf())?;
let s = AffineTransform::from_file(file.path().to_path_buf())?;
assert_eq!(s, t);
Ok(())
}
macro_rules! interp_tests_bspline {
($($name:ident: $t:ty $(,)?)*) => {
$(
#[test]
fn $name() -> Result<()> {
let j = julia_image(-120f32, 10f32)?.mapv(|x| x as $t);
let k = julia_image(0f32, 0f32)?.mapv(|x| x as $t);
let shape = j.shape();
let origin = [
((shape[1] - 1) as f64) / 2f64,
((shape[0] - 1) as f64) / 2f64,
];
let transform = AffineTransform::new([1., 0., 0., 1., 120., -10.], origin, [shape[0], shape[1]]);
let n: Array2<_> = transform.transform_image_bspline(j.view())?.into();
let d = (k.mapv(|x| x as f64) - n.mapv(|x| x as f64)).powi(2).sum();
assert!(d <= (shape[0] * shape[1]) as f64);
Ok(())
}
)*
}
}
interp_tests_bspline! {
interpbs_u8: u8,
interpbs_i8: i8,
interpbs_u16: u16,
interpbs_i16: i16,
interpbs_u32: u32,
interpbs_i32: i32,
interpbs_u64: u64,
interpbs_i64: i64,
interpbs_f32: f32,
interpbs_f64: f64,
}
macro_rules! interp_tests_nearest_neighbor {
($($name:ident: $t:ty $(,)?)*) => {
$(
#[test]
fn $name() -> Result<()> {
let j = julia_image(-120f32, 10f32)?.mapv(|x| x as $t);
let k = julia_image(0f32, 0f32)?.mapv(|x| x as $t);
let shape = j.shape();
let origin = [
((shape[1] - 1) as f64) / 2f64,
((shape[0] - 1) as f64) / 2f64,
];
let j0 = j.clone();
let k0 = k.clone();
let transform = AffineTransform::new([1., 0., 0., 1., 120., -10.], origin, [shape[0], shape[1]]);
// make sure j & k weren't mutated
assert!(j.iter().zip(j0.iter()).map(|(a, b)| a == b).all(|x| x));
assert!(k.iter().zip(k0.iter()).map(|(a, b)| a == b).all(|x| x));
let n: Array2<_> = transform.transform_image_nearest_neighbor(j.view())?.into();
let d = (k.mapv(|x| x as f64) - n.mapv(|x| x as f64)).powi(2).sum();
assert!(d <= (shape[0] * shape[1]) as f64);
Ok(())
}
)*
}
}
interp_tests_nearest_neighbor! {
interpnn_u8: u8,
interpnn_i8: i8,
interpnn_u16: u16,
interpnn_i16: i16,
interpnn_u32: u32,
interpnn_i32: i32,
interpnn_u64: u64,
interpnn_i64: i64,
interpnn_f32: f32,
interpnn_f64: f64,
}
macro_rules! registration_tests_translation {
($($name:ident: $t:ty $(,)?)*) => {
$(
#[test]
fn $name() -> Result<()> {
let j = julia_image(0f32, 0f32)?.mapv(|x| x as $t);
let k = julia_image(10f32, 20f32)?.mapv(|x| x as $t);
let j0 = j.clone();
let k0 = k.clone();
let t = AffineTransform::register_translation(j.view(), k.view())?;
// make sure j & k weren't mutated
assert!(j.iter().zip(j0.iter()).map(|(a, b)| a == b).all(|x| x));
assert!(k.iter().zip(k0.iter()).map(|(a, b)| a == b).all(|x| x));
let mut m = Array2::eye(3);
m[[0, 2]] = -10f64;
m[[1, 2]] = -20f64;
let d = (t.matrix() - m).powi(2).sum();
assert!(d < 0.01, "d: {}, t: {:?}", d, t.parameters);
Ok(())
}
)*
}
}
registration_tests_translation! {
registration_translation_u8: u8,
registration_translation_i8: i8,
registration_translation_u16: u16,
registration_translation_i16: i16,
registration_translation_u32: u32,
registration_translation_i32: i32,
registration_translation_u64: u64,
registration_translation_i64: i64,
registration_translation_f32: f32,
registration_translation_f64: f64,
}
macro_rules! registration_tests_affine {
($($name:ident: $t:ty $(,)?)*) => {
$(
#[test]
fn $name() -> Result<()> {
let j = julia_image(0f32, 0f32)?.mapv(|x| x as $t);
let shape = j.shape();
let origin = [
((shape[1] - 1) as f64) / 2f64,
((shape[0] - 1) as f64) / 2f64,
];
let s = AffineTransform::new([1.2, 0., 0., 1., 5., 7.], origin, [shape[0], shape[1]]);
let k: Array2<_> = s.transform_image_bspline(j.view())?.into();
let t = AffineTransform::register_affine(j.view(), k.view())?.inverse()?;
let d = (t.matrix() - s.matrix()).powi(2).sum();
assert!(d < 0.025, "d: {}, t: {:?}", d, t.parameters);
Ok(())
}
)*
}
}
registration_tests_affine! {
registration_tests_affine_u8: u8,
registration_tests_affine_i8: i8,
registration_tests_affine_u16: u16,
registration_tests_affine_i16: i16,
registration_tests_affine_u32: u32,
registration_tests_affine_i32: i32,
registration_tests_affine_u64: u64,
registration_tests_affine_i64: i64,
registration_tests_affine_f32: f32,
registration_tests_affine_f64: f64,
}
}