dipy_registration.py

""" Run non-linear registration using Dipy
"""

import os
from os.path import split as psplit, join as pjoin, splitext, exists
import pickle

import numpy as np
import numpy.linalg as npl

import nibabel as nib
from nibabel.affines import to_matvec

from scipy.ndimage import affine_transform

from dipy.align.imaffine import (MutualInformationMetric, AffineRegistration)
from dipy.align.transforms import (TranslationTransform3D,
                                   RigidTransform3D,
                                   AffineTransform3D)
from dipy.align.imwarp import SymmetricDiffeomorphicRegistration
from dipy.align.metrics import CCMetric


TEMPLATE_IMG = 'mni_icbm152_t1_tal_nlin_asym_09a.nii'
TEMPLATE_MASK = 'mni_icbm152_t1_tal_nlin_asym_09a_mask.nii'


def as_image(image):
    """ If `image` is string, assume filename and load image, else pass through
    """
    if isinstance(image, str):
        image = nib.load(image)
    return image


def apply_mask(brain_img, mask_img):
    """ Load brain image, matching mask; apply and return masked image

    Parameters
    ----------
    brain_img : str or image
        string giving image filename or image object
    mask_img : str or image
        string giving mask image filename or image object.  Mask must match
        `brain_img` after mapping with image and mask affine.

    Returns
    -------
    mask_img : image object
        image object where data is data from `brain_img` multiplied elementwise
        by the data of `mask_img`, where the mask data has[ been resampled into
        the voxel space of `brain_img` if necessary.
    """
    brain_img = as_image(brain_img)
    b_aff = brain_img.affine
    b_hdr = brain_img.header
    mask_img = as_image(mask_img)
    brain_data = brain_img.get_data()
    mask_data = mask_img.get_data()
    if not np.allclose(b_aff, mask_img.affine):
        # Mask and brain have different affines - we need to resample
        brain2mask = npl.inv(mask_img.affine) @ brain_img.affine
        mat, vec = to_matvec(brain2mask)
        mask_data = affine_transform(mask_data, mat, vec,
                                     output_shape=brain_data.shape,
                                     order=0)  # nearest neighbor
    return nib.Nifti1Image(brain_data * mask_data, b_aff, b_hdr)


def register_affine(t_masked, m_masked, affreg=None,
                    final_iters=(10000, 1000, 100)):
    """ Run affine registration between images `t_masked`, `m_masked`

    Parameters
    ----------
    t_masked : image
        Template image object, with image data masked to set out-of-brain
        voxels to zero.
    m_masked : image
        Moving (individual) image object, with image data masked to set
        out-of-brain voxels to zero.
    affreg : None or AffineRegistration instance, optional
        AffineRegistration with which to register `m_masked` to `t_masked`.  If
        None, we make an instance with default parameters.
    final_iters : tuple, optional
        Length 3 tuple of level iterations to use on final affine pass of the
        registration.

    Returns
    -------
    affine : shape (4, 4) ndarray
        Final affine mapping from voxels in `t_masked` to voxels in `m_masked`.
    """
    if affreg is None:
        metric = MutualInformationMetric(nbins=32,
                                         sampling_proportion=None)
        affreg = AffineRegistration(metric=metric)
    t_data = t_masked.get_data().astype(float)
    m_data = m_masked.get_data().astype(float)
    t_aff = t_masked.affine
    m_aff = m_masked.affine
    translation = affreg.optimize(t_data, m_data, TranslationTransform3D(),
                                  None, t_aff, m_aff)
    rigid = affreg.optimize(t_data, m_data, RigidTransform3D(), None, t_aff,
                            m_aff, starting_affine=translation.affine)
    # Maybe bump up iterations for last step
    if final_iters is not None:
        affreg.level_iters = list(final_iters)
    affine = affreg.optimize(t_data, m_data, AffineTransform3D(), None, t_aff,
                             m_aff, starting_affine=rigid.affine)
    return affine.affine


def register_diffeo(t_masked, m_masked, start_affine, registration=None):
    """ Run non-linear registration between `t_masked` and `m_masked`

    Parameters
    ----------
    t_masked : image
        Template image object, with image data masked to set out-of-brain
        voxels to zero.
    m_masked : image
        Moving (individual) image object, with image data masked to set
        out-of-brain voxels to zero.
    start_affine : shape (4, 4) ndarray
        Affine mapping from voxels in `t_masked` to voxels in `m_masked`.
    registration : None or SymmetricDiffeoMorphicRegistration instance
        Registration instance we will use to register `t_masked` and
        `m_masked`.  If None, make a default registration object.

    Returns
    -------
    mapping : mapping instance
        Instance giving affine + non-linear mapping between voxels in
        `t_masked` and voxels in `m_masked`.
    """
    if registration is None:
        registration = SymmetricDiffeomorphicRegistration(
            metric=CCMetric(3),
            level_iters=[10, 10, 5])
    return registration.optimize(t_masked.get_data().astype(float),
                                 m_masked.get_data().astype(float),
                                 t_masked.affine,
                                 m_masked.affine,
                                 start_affine)


def register_save(template_fname, template_mask_fname,
                  moving_fname, moving_mask_fname):
    """ Resister individual image `moving_fname` to template `template_fname`

    Save warped image and mapping object to directory of `moving_fname`.

    Parameters
    ----------
    template_fname : str
        string giving image filename of template image to register to
    template_mask_fname : str
        string giving image filename of brain mask corresponding to
        `template_fname`
    moving_fname : str
        string giving image filename of individual image to register
    moving_mask_fname : str
        string giving image filename of brain mask corresponding to
        `moving_fname`

    Returns
    -------
    mapping : mapping instance
        Instance giving affine + non-linear mapping between voxels in
        `template_fname` and voxels in `moving_fname`.
    """
    path, basename = psplit(moving_fname)
    root, ext = splitext(basename)
    if ext == '.gz':  # ignore .gz suffix
        root, ext = splitext(root)
    t_masked = apply_mask(template_fname, template_mask_fname)
    m_masked = apply_mask(moving_fname, moving_mask_fname)
    affine = register_affine(t_masked, m_masked)
    mapping = register_diffeo(t_masked, m_masked, affine)
    masked_data = m_masked.get_data()
    warped_moving = nib.Nifti1Image(mapping.transform(masked_data),
                                    t_masked.affine,
                                    t_masked.header)
    nib.save(warped_moving, pjoin(path, 'w_' + basename))
    with open(pjoin(path, 'map_' + root + '.pkl'), 'wb') as fobj:
        pickle.dump(mapping, fobj)
    return mapping


def as_mapping(mapping):
    """ If `mapping` is string, assume filename, load pickle, else pass through
    """
    if isinstance(mapping, ):
        with open(mapping, 'rb') as fobj:
            mapping = pickle.load(fobj)
    return mapping


def write_warped(fname, mapping, interpolation='nearest', template_header=None):
    """ Warp an image in individual space to template space

    Parameters
    ----------
    fmame : str
        Filename of image to warp in template space
    mapping : mapping instance
        object containing mapping from individual space to template space
    interpolation : str, optional
        interpolation to use when resampling data from `fname`
    template_header : None or header instance
        template header with which to save image.  If None, use default header.
    """
    img = nib.load(fname)
    mapping = as_mapping(mapping)
    template_affine = mapping.codomain_grid2world
    data = img.get_data().astype(float)
    warped = mapping.transform(data, interpolation=interpolation)
    warped_img = nib.Nifti1Image(warped, template_affine, template_header)
    path, basename = psplit(fname)
    out_fname = pjoin(path, 'w_' + basename)
    nib.save(warped_img, out_fname)


def find_anatomicals(root):
    """ Find anatomical image, mask pairs from OpenFMRI directory root `root`

    Parameters
    ----------
    root : str
        root directory of OpenFMRI dataset - e.g. "ds114"

    Returns
    -------
    anatomicals : list
        List of anatomical, mask image pairs for each subject in `root`
    """
    anatomicals = []
    for dirpath, dirnames, filenames in os.walk(root):
        if not 'highres001.nii.gz' in filenames:
            continue
        full_image = pjoin(dirpath, 'highres001.nii.gz')
        mask_image = pjoin(dirpath, 'highres001_brain_mask.nii.gz')
        assert exists(full_image)
        assert exists(mask_image)
        anatomicals.append((mask_image, full_image))
    return anatomicals


def sub2img_mask(root, sub_no):
    """ Return anatomical image, mask pair from OpenFMRI root, subject no

    Parameters
    ----------
    root : str
        root directory of OpenFMRI dataset - e.g. "ds114"
    sub_no : int
        subject no, where 1 is the first subject

    Returns
    -------
    anatomical_fname : str
        filename of anatomical image, beginning with path `root`
    mask_fname : str
        filename of anatomical image brain mask, beginning with path `root`
    """
    anatomical_path = pjoin(root, 'sub{:03d}'.format(sub_no), 'anatomy')
    ret = (pjoin(anatomical_path, 'highres001.nii.gz'),
           pjoin(anatomical_path, 'highres001_brain_mask.nii.gz'))
    if all(exists(p) for p in ret):
        return ret
    return ()


def write_highres(root):
    """ Calculcate parameters, write anatomicals from OpenFMRI directory `root`

    Parameters
    ----------
    root : str
        root directory of OpenFMRI dataset - e.g. "ds114"
    """
    for moving_img, moving_mask in find_anatomicals(root):
        register_save(TEMPLATE_IMG, TEMPLATE_MASK,
                      moving_img, moving_mask)


def write_highres_parallel(root):
    """ Calculcate parameters, write anatomicals from OpenFMRI directory `root`

    Use parallel execution.  Careful, this can crash your machine with too many
    images found at `root`.

    Parameters
    ----------
    root : str
        root directory of OpenFMRI dataset - e.g. "ds114"
    """
    import multiprocessing
    jobs = []
    for moving_img, moving_mask in find_anatomicals(root):
        p = multiprocessing.Process(target=register_save, args=(
            TEMPLATE_IMG, TEMPLATE_MASK, moving_img, moving_mask))
        jobs.append(p)
        p.start()


def register_subject(root, sub_no):
    """ Calculcate parameters, write anatomical for subject at OpenFMRI `root`

    Parameters
    ----------
    root : str
        root directory of OpenFMRI dataset - e.g. "ds114"
    sub_no : int
        subject no, where 1 is the first subject

    Returns
    -------
    mapping : mapping instance
        Instance giving affine + non-linear mapping between voxels in
        `t_masked` and voxels in `m_masked`.
    """
    moving_img, moving_mask = sub2img_mask(root, sub_no)
    return register_save(TEMPLATE_IMG, TEMPLATE_MASK,
                         moving_img, moving_mask)