[BUG] MRI-CT alignment failure

[alexrockhill]

Description

It looks like there is something suboptimal/unstable about the way Dipy does the image registration relative to antspyx. Hopefully it's just a small fix. Here is how to reproduce the problem.

CT file: https://drive.google.com/file/d/1Epaw_kVLElBIbgo184KIsLSH7YuGliLI/view?usp=sharing MR file: https://drive.google.com/file/d/1deiJMNH00NU0z3Lr6IPUe0HI7EKjyTo-/view?usp=sharing

import numpy as np
from dipy.align import affine_registration, center_of_mass, translation, rigid
import nibabel as nib
import matplotlib.pyplot as plt

def plot_overlay(image, compare):
    """Define a helper function for comparing plots."""
    image = nib.orientations.apply_orientation(
        np.asarray(image.dataobj), nib.orientations.axcodes2ornt(
            nib.orientations.aff2axcodes(image.affine))).astype(np.float32)
    compare = nib.orientations.apply_orientation(
        np.asarray(compare.dataobj), nib.orientations.axcodes2ornt(
            nib.orientations.aff2axcodes(compare.affine))).astype(np.float32)
    fig, axes = plt.subplots(1, 3, figsize=(12, 4))
    for i, ax in enumerate(axes):
        ax.imshow(np.take(image, [image.shape[i] // 2], axis=i).squeeze().T,
                  cmap='gray')
        ax.imshow(np.take(compare, [compare.shape[i] // 2],
                          axis=i).squeeze().T, cmap='gist_heat', alpha=0.5)
        ax.invert_yaxis()
        ax.axis('off')
    fig.tight_layout()

moving = nib.load('sub-12_CT.nii.gz')
static = nib.load('sub-12_T1.mgz')

_, reg_affine = affine_registration(
    np.array(moving.dataobj),
    np.array(static.dataobj),
    moving_affine=moving.affine,
    static_affine=static.affine,
    nbins=32,
    metric='MI',
    pipeline=[center_of_mass, translation])

moved_img, reg_affine = affine_registration(
    np.array(moving.dataobj),
    np.array(static.dataobj),
    moving_affine=moving.affine,
    static_affine=static.affine,
    nbins=32,
    metric='MI',
    pipeline=[rigid],
    starting_affine=reg_affine)

moved = nib.MGHImage(moved_img.astype(np.float32), static.affine)
plot_overlay(static, moved)

CT_aligned = mne.transforms.apply_volume_registration(CT_orig, T1, reg_affine)

# compare to ANTS
import ants
moving = ants.image_read('sub-12_CT.nii.gz')
fixed = ants.image_read('sub-12_T1.mgz')

trans = ants.registration(fixed=fixed, moving=moving, type_of_transform='Rigid')
ants.image_write(trans['warpedmovout'], 'sub-12_CT_aligned_ANTS.mgz')
moved = nib.load('sub-12_CT_aligned_ANTS.mgz')

plot_overlay(static, moved)

This is on a Mac running Catalina with Dipy version 1.5.0dev using Python 3.9.4. cc @larsoner @leosunpsy

Related to https://github.com/mne-tools/mne-python/issues/10011

Way to reproduce

[If reporting a bug, please include the following important information:]

• [x] Code example
• [x] Relevant images (if any)
• [x] Operating system and version (run python -c "import platform; print(platform.platform())")
• [x] Python version (run python -c "import sys; print("Python", sys.version)")
• [x] dipy version (run python -c "import dipy; print(dipy.__version__)")
• [x] dependency version (numpy, scipy, nibabel, h5py, cvxpy, fury)
  • import numpy; print("NumPy", numpy.version)
  • import scipy; print("SciPy", scipy.version)
  • import nibabel; print("Nibabel", nibabel.version)
  • import h5py; print("H5py", h5py.version)
  • import cvxpy; print("Cvxpy", cvxpy.version)
  • import fury; print("fury", fury.version)

[alexrockhill]

Ok, here's what I got from ANTS

And here is with the above code in Dipy

Hopefully you can see it's at least several mm off, it is easier to see in freeview in 3D.

[alexrockhill]

Update: interestingly ants version 0.3.1 gives the same poor alignment using a Mac running Catalina but on an Ubuntu 18.04 with ants version 0.2.9 the alignment is correct (shown above). I just checked 0.2.9 on Mac and it also fails to align for ants. That's a bit difficult as this leaves Ubuntu as the only OS I've been able to get a good alignment on and that the setup depends on the OS does not sound good. The Mac has numpy version 1.20.3 whereas the Ubuntu (working) has 1.23.3 so that may be the issue as well.

Update: the ants version doesn't work with numpy 1.21.4 either on Mac. I'm going to stop troubleshooting here as my best guess is that something is going weird with the OS on ants but that's not really the topic of discussion here. The main point is that there is a good registration and it is attainable but in Dipy it isn't achieved with the latest version. If anyone has any thoughts on this, that would be very much appreciated.

[skoudoro]

Thank you for this nice report!

this leaves Ubuntu as the only OS....

it is funny that #2376 says the opposite. Good result on Mac and bad on Linux. It seems there is no consensus.

Edit: I have a doubt finally on this old discussion, might be the same conclusion.

the setup depends on the OS does not sound good.

Agree it does not sound good to have such a difference between os. I wonder if our wheels builder compiler could introduce floating issues (gcc vs clang)

Overall, we need to investigate RigidTransform3D. There is clearly something wrong that we do not catch during our test. I do not think it will be an easy task but we will keep you updated. Also, I don't think the NumPy version has an impact on this.

Let us know also if you have some ideas or complementary information.

[alexrockhill]

Thanks for getting back to me @skoudoro! I'm happy to do my part/help as much as I can get this fixed. I think this data is a good benchmark as Dipy worked great for 6/8 subjects in a recent coregistration pipeline but this was one of the ones that failed. Not sure what about it is such a challenge but maybe the bounding boxes are especially off as you can see in the correct alignment part of the skull is cut off in the CT at the bottom.

To clarify, I only got a correct alignment with this data using ANTS and not Dipy which I included since Dipy uses the same algorithm (ASAIK) so potentially ANTS might have a difference that might fix the issue.

I've been running a few more tests but this is definitely a tricky one as each registration takes ~15 or 20 minutes to run. Maybe once the problem is pinpointed a more directed test can be added though.

I really don't know much about the compilers and implementation but I'm happy to look into it a bit more and try and get this fixed. Thanks again :)

[alexrockhill]

One last point is that after only translation it looks like this

So I'm guessing the initial alignment is causing the algorithm to get stuck in a local minimum when it goes to optimize the rigid step... Not exactly sure how to fix this though. In MNE, we wrap the function to zoom the images first so I can check if zooming the translation and center of mass steps might help set up the rigid step better.

[alexrockhill]

Here is the correctly aligned CT image https://drive.google.com/file/d/1ZYV6tY9GbcyFrMd4FN2ecFKoPKzpWB1J/view?usp=sharing

When I do the following:

from dipy.align.imaffine import MutualInformationMetric
from dipy.align.transforms import RigidTransform3D
import nibabel as nib

metric = MutualInformationMetric()
transform = RigidTransform3D()
# used `static` and `moved` from above
metric.setup(transform, np.array(static.dataobj), np.array(moved.dataobj),
    static.affine, moved.affine)
metric._update_mutual_information(transform.get_identity_parameters())
print(f'Suboptimal "Dipy main" metric: {metric.metric_val}')

moved_good = nib.load('sub-12_CT_aligned_good.mgz')
metric.setup(transform, np.array(static.dataobj), np.array(moved_good.dataobj),
    static.affine, moved_good.affine)
metric._update_mutual_information(transform.get_identity_parameters())
print(f'Ideal alignment (ANTS on Ubuntu) metric: {metric.metric_val}')

# sanity check, original unaligned metric
metric.setup(transform, np.array(static.dataobj), np.array(moving.dataobj),
    static.affine, moving.affine)
metric._update_mutual_information(transform.get_identity_parameters())
print(f'Original unaligned metric: {metric.metric_val}')

I get

Suboptimal "Dipy main" metric: 0.10511911041182931
Ideal alignment (ANTS on Ubuntu) metric: 0.13006035331117444
Original unaligned metric: 0.01899945889571585

How I interpret this is that the ideal alignment does indeed have a better score on the metric but the optimizer is not getting there. Please let me know if this is incorrect. Perhaps then, I could write a PR to add basinhopping as the optimizer instead of minimize (with a keyword argument). It looks like there are some arguments as well as params0 passed to minimize so maybe this is not as straightforward as I'd hope... I'll hold off on doing this unless @skoudoro or others think this is a good idea.

[alexrockhill]

To update, basinhopping even from the slightly incorrect alignment is computationally prohibitive to the point of being unfeasible. I ran 100 iterations for 10+ hours before killing and now with 5 iterations, it's still running after an hour. Even if it does work, I don't think it's a practical solution, at least not without fine tuning. I'm really not sure where to go from here. Maybe @larsoner if you had a second to weigh in that would be amazing: tl;dr it seems the affine optimization is getting stuck in a local minimum using scipy.optimize.minimize. Maybe a method other than L-BFGS-B would help?

[alexrockhill]

Okay I tried Nelder-Mead as the solver and it was no better. I'm a bit out of ideas here but I'll think on it.

import numpy as np
import nibabel as nib
from dipy.align import affine_registration
from dipy.align.imaffine import AffineRegistration
from dipy.align.transforms import RigidTransform3D

moving = nib.load('sub-12_CT.nii.gz')
static = nib.load('sub-12_T1.mgz')

moving2, reg_affine = affine_registration(
    np.array(moving.dataobj),
    np.array(static.dataobj),
    moving_affine=moving.affine,
    static_affine=static.affine,
    nbins=32,
    metric='MI',
    pipeline=[center_of_mass, translation])

moving2 = nib.MGHImage(moving2.astype(np.float32), static.affine)

affreg = AffineRegistration(method='Nelder-Mead')
trans = affreg.optimize(np.array(static.dataobj), np.array(moving2.dataobj),
    RigidTransform3D(), None, static.affine, moving2.affine)

moved2 = nib.MGHImage(trans.transform(np.array(moving.dataobj)).astype(np.float32), static.affine)

[alexrockhill]

Hmmm I'm a bit stuck on this, any thoughts of how to fix it?

[Garyfallidis]

@alexrockhill will get back to you on this asap. We are currently working to fix another issue that blocks the current release. I see that @larsoner is planning to try a few ideas too. Which is great. More asap.

[Garyfallidis]

Also thank you for the detailed analysis @alexrockhill .

[larsoner]

ANTS result varies

I made a conda env and pip installed antspyx-0.3.1, and if I run this snippet:

``` import numpy as np from dipy.align.imaffine import MutualInformationMetric from dipy.align.transforms import RigidTransform3D import nibabel as nib moving = nib.load('sub-12_CT.nii.gz') static = nib.load('sub-12_T1.mgz') transform = RigidTransform3D() metric = MutualInformationMetric() import ants mov = ants.image_read('sub-12_CT.nii.gz') fixed = ants.image_read('sub-12_T1.mgz') trans = ants.registration(fixed=fixed, moving=mov, type_of_transform='Rigid') ants.image_write(trans['warpedmovout'], 'sub-12_CT_aligned_ANTS.mgz') moved_good = nib.load('sub-12_CT_aligned_ANTS.mgz') metric.setup(transform, np.array(static.dataobj), np.array(moved_good.dataobj), static.affine, moved_good.affine) metric._update_mutual_information(transform.get_identity_parameters()) print(f'Ideal alignment (ANTS on Ubuntu) metric: {metric.metric_val}') ```

I get a different, worse value from what you got @alexrockhill that is more in line with what dipy produces, and it changes when I run my same python check.py script multiple times (different outputs combined below):

Ideal alignment (ANTS on Ubuntu) metric: 0.11514645815260521
Ideal alignment (ANTS on Ubuntu) metric: 0.11716000206390613
Ideal alignment (ANTS on Ubuntu) metric: 0.11588613778150773

Can you try running that code snippet in a fresh directory with the binaries you linked to make sure you get a different MI with the exact same inputs and code, and see if you see similar variability? To me it suggests there is some non-deterministic / randomization business going on in ANTS that could be affecting our comparisons.

Cannot replicate dipy result

Using the ANTS image that you produced @alexrockhill but redoing the dipy computations on my Ubuntu machine I get:

Suboptimal "Dipy main" metric: 0.1147438118990917
Ideal alignment (ANTS on Ubuntu) metric: 0.13006035331117444
Original unaligned metric: 0.01899945889571585

This doesn't match your value of 0.1051.

Testing optimization methods

I created a little branch that allows playing with the optimization a bit in https://github.com/dipy/dipy/pull/2498. When using a self-contained adaptation of your code (assuming you download the three files linked above) then I do see this in the more verbose output when I pass options=dict(iprint=1, gtol=0) for example:

 Warning:  more than 10 function and gradient
   evaluations in the last line search.  Termination
   may possibly be caused by a bad search direction.

Maybe this is indicative of some problem? Not sure.

For what it's worth, in addition to the default l-bfgs-b (0.1135), I tried newton-cg (0.1135), bfgs (0.1106), slsqp (0.1116), nelder-mead (0.1020), and cobyla (0.1077), all of which gave similar (or worse) results. slsqp converges much faster to a pretty similar solution -- @alexrockhill you might try some global optimizer like basinhopping using slsqp for local minimization... maybe it'll be fast enough to get out of the local minimum that dipy seems to get stuck in?

One final idea -- @alexrockhill have you checked that taking the affine transformation matrix that ANTS computes and applying that using dipy's resampling functions gives an identical image (or at least nearly the same mutual information) as you get with the sub-12_CT_aligned_good.mgz image that was written by ANTS? We should rule out that the interpolation parameters (nearest, linear, cubic) is somehow affecting the comparison via mutual information here.

[alexrockhill]

Thanks for working on this @larsoner!

I ran the code on my Ubuntu machine (ANTS version 0.2.9) and I get: Ideal alignment (ANTS on Ubuntu) metric: 0.11627779721091346 which I thought I would get around 0.13 but replicates what you got. This was suboptimal (I checked in Freeview) so now the only working alignment I got isn't working anymore... It might be stochastic though, let me run it two more times and check.

EDIT: I got Ideal alignment (ANTS on Ubuntu) metric: 0.11487923730914311 and Ideal alignment (ANTS on Ubuntu) metric: 0.11540763572240431 for another two runs so I can confirm that there is non-seeded randomness.

That's an interesting idea of basinhopping with a different optimizer, I can give that a try. I'm imagining a function like dipy.align.affine_registration2 or dipy.align.affine_registration_fix which could be run after the first fit to get a global optimum.

I did check with other subjects that the interpolation parameters don't affect anything. It's rather odd, I could not for the life of me figure out how to get an affine matrix out of the metadata in the ANTS alignment so I just registered the original CT to ANTS-aligned "good" CT using dipy.align.affine_registration and that gave an affine matrix that worked just fine.

[drombas]

Hi! I checked the optimized metric returned by affine_registration and it doesn't match the metric computed separately afterwards. Not sure if this is expected or could indicate something is wrong with the metric (I think the optimizer minimizes -MutualInformation so I don't see why the values differ).

metric = MutualInformationMetric()
transform = RigidTransform3D()

static = nib.load('sub-12_T1.mgz')
moving = nib.load('sub-12_CT.nii.gz')

moved, aff_reg, xopt, fopt = affine_registration(
    moving=moving.get_fdata(),
    static=static.get_fdata(),
    moving_affine=moving.affine,
    static_affine=static.affine,
    pipeline=["center_of_mass","translation","rigid"],
    ret_metric=True)

metric.setup(transform, static.get_fdata(), moved, static.affine, static.affine)
metric._update_mutual_information(transform.get_identity_parameters())
print(f'Dipy (computed): {metric.metric_val}')
print(f'Dipy (optimized): {-fopt}')

Dipy (computed): 0.11474381189837092
Dipy (optimized): 0.32778350621531893

Note: in the latest version the pipeline argument of affine_registration is a list of strings instead of functions but I don't think that should matter

[alexrockhill]

As another update, basinhopping with slsqp as the optimizer and niter=5 is still running after 20 minutes. Since the amount of time is greater than the original algorithm, I don't think this will work (and it's not done yet). It's possible that we could write our own "basinhopping" fairly quickly, which would just involve taking several random starting points and using the best optimization.

Otherwise, I think manually aligning and then running the normal dipy.align.affine_registration works with the obvious downside that involves time spent manually aligning.

[alexrockhill]

For help troubleshooting, the "Dipy master suboptimal alignment" affine is

array([[   0.99998698,   -0.00262631,    0.00437602, -110.3154669 ],
       [   0.00264113,    0.99999078,   -0.00338521,  -85.13388839],
       [  -0.00436709,    0.00339672,    0.9999847 , -105.7693801 ],
       [   0.        ,    0.        ,    0.        ,    1.        ]])

And the Dipy alignment of the original CT to the good CT as found by ANTS is

array([[   0.99897466,   -0.03035359,    0.03359012, -110.73023661],
       [   0.03147058,    0.9989518 ,   -0.0332402 ,  -84.27371857],
       [  -0.03254595,    0.03426322,    0.99888277, -105.72314593],
       [   0.        ,    0.        ,    0.        ,    1.        ]])

[larsoner]

So I think the TL;DR is that both ANTS and dipy can get stuck in a local minimum for this particular problem. I'm not sure there is much to do at the dipy end about this, other than note in the docstring somewhere that this can happen, and suggest that people check alignments and use manual alignment in errant cases?

[alexrockhill]

Ah after thinking about it, I think I have an idea. If the alignment fails, you could divide the image up into 8 sub-images and fit all of those and take the best MI metric. This could be done with existing Dipy functions without modification. I'll give it a try for this example.

import numpy as np
from dipy.align import affine_registration, center_of_mass, translation, rigid, resample
from dipy.align.imaffine import MutualInformationMetric
from dipy.align.transforms import RigidTransform3D
import nibabel as nib
from itertools import product

moving = nib.load('sub-12_CT.nii.gz')
static = nib.load('sub-12_T1.mgz')

moving_data = np.array(moving.dataobj).astype(np.float32)
static_data = np.array(static.dataobj).astype(np.float32)

_, reg_affine = affine_registration(
    moving_data,
    static_data,
    moving_affine=moving.affine,
    static_affine=static.affine,
    nbins=32,
    metric='MI',
    pipeline=[center_of_mass, translation])

fits = dict()
angle = np.deg2rad(10)
x_rot = lambda theta: np.array([[1, 0, 0, 0], [0, np.cos(theta), -np.sin(theta), 0],
                                                      [0, np.sin(theta), np.cos(theta), 0], [0, 0, 0, 1]], dtype=float)
y_rot = lambda theta: np.array([[np.cos(theta), 0, np.sin(theta), 0], [0, 1, 0, 0], 
                                                      [-np.sin(theta), 0, np.cos(theta), 0], [0, 0, 0, 1]], dtype=float)
z_rot = lambda theta: np.array([[np.cos(theta), -np.sin(theta), 0, 0],
                                                      [np.sin(theta), np.cos(theta), 0, 0], [0, 0, 1, 0], [0, 0, 0, 1]], dtype=float)
for factors in product(*[range(-1, 2)] * 3):
    rot = x_rot(factors[0] * angle).dot(y_rot(factors[1] * angle)).dot(z_rot(factors[2] * angle))
    _, rigid_reg_affine = affine_registration(
        moving_data,
        static_data,
        moving_affine=moving.affine,
        static_affine=static.affine,
        nbins=32,
        metric='MI',
        pipeline=[rigid],
        starting_affine=np.dot(reg_affine, rot))
    moved = resample(
        moving=moving_data,
        moving_affine=moving.affine,
        static=static_data,
        static_affine=static.affine,
        between_affine=rigid_reg_affine)
    metric = MutualInformationMetric()
    transform = RigidTransform3D()
    metric.setup(transform, static_data, np.array(moved.dataobj), static.affine, static.affine)
    metric._update_mutual_information(transform.get_identity_parameters())
    fits[metric.metric_val] = rigid_reg_affine

print(sorted(fits.keys()))
rigid_reg_affine = fits[max(fits])

moved = resample(
        moving=moving_data,
        moving_affine=moving.affine,
        static=static_data,
        static_affine=static.affine,
        between_affine=rigid_reg_affine)
nib.save(moved, 'sub-12_CT_aligned_best.mgz')

That's what I was thinking, I'll run it and update when I get the results.

EDIT: I'm trying one more idea by modifying this example that was subdividing the image to instead be starting from different rotations.

[alexrockhill]

Nope, didn't work here's what was printed

[0.04834651995765701, 0.06616670645292315, 0.06938223382793579, 0.09937423106217762, 0.10319763311401721, 0.10880012822782431, 0.11236077393632106, 0.11291619515587295]

I'm out of ideas, I think they'll just have to be manually aligned.

I'll close this tomorrow unless anyone else has ideas.

[alexrockhill]

I tried another idea of starting with different initial rotations (above) and no dice for that either. I'm going to go ahead and close and say there isn't a good solution at this time.