align.py

import numpy as np
from dipy.align import affine_registration, center_of_mass, translation, rigid
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')

reg_kwargs = dict(
    moving_affine=moving.affine,
    static_affine=static.affine,
    nbins=32,
    metric='MI',
    #method='cobyla',
    #options=dict(rhobeg=1e-4),
    options=dict(iprint=2, gtol=0),
)
_, reg_affine = affine_registration(
    np.array(moving.dataobj), np.array(static.dataobj),
    pipeline=[center_of_mass, translation], **reg_kwargs)

moved_img, reg_affine = affine_registration(
    np.array(moving.dataobj), np.array(static.dataobj),
    pipeline=[rigid], starting_affine=reg_affine, **reg_kwargs)


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

metric = MutualInformationMetric()
transform = RigidTransform3D()
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}')