motion correction for dwi

motion_correct.py

"""
==========================================
Motion correction of DWI data
==========================================
"""

import numpy as np
from dipy.viz import regtools
from dipy.data import fetch_stanford_hardi, read_stanford_hardi
from dipy.data.fetcher import fetch_syn_data, read_syn_data
from dipy.align.imaffine import (transform_centers_of_mass,
                                 AffineMap,
                                 MutualInformationMetric,
                                 AffineRegistration)
from dipy.align.transforms import (TranslationTransform3D,
                                   RigidTransform3D,
                                   AffineTransform3D)

from dipy.io.image import load_nifti, save_nifti
from dipy.io.gradients import read_bvals_bvecs
from dipy.core.gradients import gradient_table, reorient_bvecs


def affine_reg(static, static_grid2world,
               moving, moving_grid2world):

    c_of_mass = transform_centers_of_mass(static,
                                          static_grid2world,
                                          moving,
                                          moving_grid2world)

    nbins = 32
    sampling_prop = None
    metric = MutualInformationMetric(nbins, sampling_prop)

    level_iters = [500, 100, 10]

    sigmas = [3.0, 1.0, 0.0]

    factors = [4, 2, 1]

    affreg = AffineRegistration(metric=metric,
                                level_iters=level_iters,
                                sigmas=sigmas,
                                factors=factors)

    transform = TranslationTransform3D()
    params0 = None
    starting_affine = c_of_mass.affine
    translation = affreg.optimize(static, moving, transform, params0,
                                  static_grid2world, moving_grid2world,
                                  starting_affine=starting_affine)

    transformed = translation.transform(moving)

    transform = RigidTransform3D()
    params0 = None
    starting_affine = translation.affine
    rigid = affreg.optimize(static, moving, transform, params0,
                            static_grid2world, moving_grid2world,
                            starting_affine=starting_affine)
    transformed = rigid.transform(moving)

    transform = AffineTransform3D()
    params0 = None
    starting_affine = rigid.affine
    affine = affreg.optimize(static, moving, transform, params0,
                             static_grid2world, moving_grid2world,
                             starting_affine=starting_affine)

    transformed = affine.transform(moving)

    affine.affine

    return transformed, affine.affine


fdwi = '20181031094549_resolve_4scan_trace_tra_p2_224_7.nii.gz'
fbval = '20181031094549_resolve_4scan_trace_tra_p2_224_7.bval'
fbvec = '20181031094549_resolve_4scan_trace_tra_p2_224_7.bvec'
fmask = 'out/brain_mask.nii.gz'

fcorr_dwi = 'out/dwi_corr.nii.gz'
fcorr_bval = 'out/corr.bval'
fcorr_bvec = 'out/corr.bvec'

data, affine = load_nifti(fdwi)
mask, _ = load_nifti(fmask)

static_grid2world = affine

moving_grid2world = affine

bvals, bvecs = read_bvals_bvecs(fbval, fbvec)

gtab = gradient_table(bvals, bvecs)

data_b0s = data[..., gtab.b0s_mask]

# ATTENTION: the template S0 should be after registering
# all b0s to one b0. Here I am only calculating an average
# S0. But you should register all S0s together.
data_avg_b0 = data.mean(axis=-1)
static = data_avg_b0

reg_affines = []

data_corr = np.zeros(data.shape)

for i in range(data.shape[-1]):
    print('Volume index %d' % (i,))
    if not gtab.b0s_mask[i]:

        print('Affine registration started')
        moving = data[..., i]
        moved, trans_affine = affine_reg(static, static_grid2world,
                                         moving, moving_grid2world)
        reg_affines.append(trans_affine)
        print('Affine registration finished')
    else:
        moved = data[..., i]
        trans_affine = affine

    data_corr[..., i] = moved

gtab_corr = reorient_bvecs(gtab, reg_affines)

save_nifti(fcorr_dwi, data_corr, affine)

np.savetxt(fcorr_bval, bvals)
np.savetxt(fcorr_bvec, gtab_corr.bvecs)