tensorimage.py

#!/usr/bin/env python
#
# tensorimage.py - The TensorImage class.
#
# Author: Paul McCarthy <pauldmccarthy@gmail.com>
#
"""This module provides the :class:`.TensorImage` class, which encapsulates
the diffusion tensor data generated by the FSL ``dtifit`` tool.
"""


import                   logging
import                   re
import                   glob
import os.path        as op

import fsl.data.image as fslimage


log = logging.getLogger(__name__)


def getTensorDataPrefix(path):
    """Returns the prefix used for the DTI file names in the given
    directory, or ``None`` if the DTI files could not be identified.
    """
    
    v1s   = glob.glob(op.join(path, '*_V1.*'))
    v2s   = glob.glob(op.join(path, '*_V2.*'))
    v3s   = glob.glob(op.join(path, '*_V3.*'))
    l1s   = glob.glob(op.join(path, '*_L1.*'))
    l2s   = glob.glob(op.join(path, '*_L2.*'))
    l3s   = glob.glob(op.join(path, '*_L3.*'))
    fas   = glob.glob(op.join(path, '*_FA.*'))
    mds   = glob.glob(op.join(path, '*_MD.*'))
    files = [v1s, v2s, v3s, l1s, l2s, l3s, fas, mds]

    # Gather all of the existing file
    # prefixes into a dictionary of
    # prefix : [file list] mappings.
    pattern  = '^(.*)_(?:V1|V2|V3|L1|L2|L3|FA|MD).*$'
    prefixes = {}

    for f in [f for flist in files for f in flist]:
        prefix = re.findall(pattern, f)[0]

        if prefix not in prefixes: prefixes[prefix] = [f]
        else:                      prefixes[prefix].append(f)
                
    # Discard any prefixes which are 
    # not present for every file type.
    for prefix, files in list(prefixes.items()):
        if len(files) != 8:
            prefixes.pop(prefix)

    # Discard any prefixes which
    # match any files that do
    # not look like image files
    for prefix, files in list(prefixes.items()):
        if not all([fslimage.looksLikeImage(f) for f in files]):
            prefixes.pop(prefix)

    prefixes = list(prefixes.keys())

    # No more prefixes remaining -
    # this is probably not a dtifit
    # directory
    if len(prefixes) == 0:
        return None

    # If there's more than one remaining
    # prefix, I don't know what to do -
    # just return the first one. 
    if len(prefixes) > 1:
        log.warning('Multiple dtifit prefixes detected: {}'.format(prefixes))

    return op.basename(prefixes[0])


def isPathToTensorData(path):
    """Returns ``True`` if the given directory path looks like it contains
    ``dtifit`` data, ``False`` otherwise.
    """
    
    return getTensorDataPrefix(path) is not None


class TensorImage(fslimage.Nifti1):
    """The ``TensorImage`` class is able to load and encapsulate the diffusion
    tensor data generated by the FSL ``dtifit`` tool.
    """

    
    def __init__(self, path):
        """Create a ``TensorImage``.

        :arg path: A path to a ``dtifit`` directory. Alternately, the ``path``
                   may be a dictionary with keys
                   ``{'v1', 'v2', 'v3', 'l1', 'l2', 'l3'}``, which specify
                   paths to images containing the tensor eigenvectors and
                   eigenvalues.
        """

        dtifitDir = isinstance(path, basestring)

        if dtifitDir:

            prefix = getTensorDataPrefix(path)

            if prefix is None:
                raise ValueError('Invalid path: {}'.format(path))

            v1 = op.join(path, '{}_V1'.format(prefix))
            v2 = op.join(path, '{}_V2'.format(prefix))
            v3 = op.join(path, '{}_V3'.format(prefix))
            l1 = op.join(path, '{}_L1'.format(prefix))
            l2 = op.join(path, '{}_L2'.format(prefix))
            l3 = op.join(path, '{}_L3'.format(prefix))

            paths = {'v1' : v1, 'v2' : v2, 'v3' : v3,
                     'l1' : l1, 'l2' : l2, 'l3' : l3}
            
        else:
            paths = path

        fslimage.Nifti1.__init__(self, paths['l1'], loadData=False)

        self.__v1 = fslimage.Image(paths['v1'])
        self.__v2 = fslimage.Image(paths['v2'])
        self.__v3 = fslimage.Image(paths['v3'])
        self.__l1 = fslimage.Image(paths['l1'])
        self.__l2 = fslimage.Image(paths['l2'])
        self.__l3 = fslimage.Image(paths['l3'])

        l1dir = op.abspath(op.dirname(paths['l1']))

        self.dataSource = l1dir
        self.name       = '{}[tensor]'.format(op.basename(l1dir))

        
    def V1(self): return self.__v1
    def V2(self): return self.__v2
    def V3(self): return self.__v3
    def L1(self): return self.__l1
    def L2(self): return self.__l2
    def L3(self): return self.__l3