Enhancements to freesurfer module to support different vertex data types.

fsl/data/freesurfer.py

@@ -7,53 +7,111 @@
 """This module provides the :class:`FreesurferMesh` class, which can be
 used for loading Freesurfer geometry and vertex data files.
 
+
+The following types of files generated by Freesurfer are recognised by this
+module:
+
+
+ - Core geometry files, defining the cortical mesh. Each core geometry
+   file has the same number of vertices and same triangle definitions.
+
+ - Extra geometry files, defining the cortical mesh. Extra geometry
+   files may have a different number of vertices and/or triangle
+   definitions.
+
+ - Vertex data files (a.k.a. ``'curv'`` files), containing a scalar value for
+   every vertex in the mesh. This data may also be contained in
+   ``mgz``/``mgh`` files.
+
+ - Label files, containing indices of a sub-set of vertices in the
+   mesh.
+
+ - Annotation files, containing a label value, and an RGBA colour for a
+   subset of vertices in the mesh.
+
+
 The following functions are also available:
 
   .. autosummary::
      :nosignatures:
 
-     loadFreesurferVertexFile
-     relatedFiles
+     loadVertexDataFile
+     isCoreGeometryFile
+     isGeometryFile
+     isVertexDataFile
+     isVertexLabelFile
+     isVertexAnnotFile
+     relatedGeometryFiles
+     relatedVertexDataFiles
+     findReferenceImage
 """
 
 
-import os.path as op
-import            glob
+import os.path   as op
+import itertools as it
+import              glob
+import              fnmatch
+import              collections
 
+import numpy              as np
 import nibabel.freesurfer as nibfs
 
 import fsl.utils.path    as fslpath
+import fsl.utils.memoize as memoize
+import fsl.data.image    as fslimage
 import fsl.data.mghimage as fslmgh
 import fsl.data.mesh     as fslmesh
 
 
-GEOMETRY_EXTENSIONS = ['.pial',
-                       '.white',
-                       '.sphere',
-                       '.inflated',
-                       '.orig',
-                       '.mid']
-"""File extensions which are interpreted as Freesurfer geometry files. """
+CORE_GEOMETRY_FILES = ['?h.orig',
+                       '?h.pial',
+                       '?h.white',
+                       '?h.inflated',
+                       '?h.sphere']
+"""File patterns for identifying the core Freesurfer geometry files. """
 
 
-EXTENSION_DESCRIPTIONS = [
-    "Freesurfer surface",
-    "Freesurfer surface",
-    "Freesurfer surface",
-    "Freesurfer surface",
-    "Freesurfer surface",
-    "Freesurfer surface"]
+CORE_GEOMETRY_DESCRIPTIONS = [
+    "Freesurfer surface (original)",
+    "Freesurfer surface (pial)",
+    "Freesurfer surface (white matter)",
+    "Freesurfer surface (inflated)",
+    "Freesurfer surface (sphere)"]
 """A description for each extension in :attr:`GEOMETRY_EXTENSIONS`. """
 
 
-VERTEX_DATA_EXTENSIONS = ['.curv',
-                          '.crv',
-                          '.area',
-                          '.thickness',
-                          '.volume',
-                          '.mgh',
-                          '.mgz']
-"""File extensions which are interpreted as Freesurfer vertex data files. """
+EXTRA_GEOMETRY_FILES = ['?h.orig.nofix',
+                        '?h.smoothwm.nofix',
+                        '?h.inflated.nofix',
+                        '?h.qsphere.nofix',
+                        '?h.sphere.nofix',
+                        '?h.white.preaparc']
+"""Other geometry files which may be present in Freesurfer output. """
+
+
+VERTEX_DATA_FILES = ['?h.thickness',
+                     '?h.curv',
+                     '?h.area',
+                     '?h.sulc',
+                     '?h.*.stats',
+                     '.mgh',
+                     '.mgz']
+"""File patterns which are interpreted as Freesurfer vertex data files,
+containing a scalar value for every vertex in the mesh.
+"""
+
+
+VERTEX_LABEL_FILES = ['?h.*.label']
+"""File patterns which are interpreted as Freesurfer vertex label files,
+containing a scalar value for a sub-set of vertices in the mesh.
+"""
+
+
+VERTEX_ANNOT_FILES = ['?h.*.annot']
+"""File patterns which are interpreted as Freesurfer vertex annotation files,
+containing a scalar value and an RGBA colour for a sub-set of vertices in the
+mesh.
+"""
 
 
 class FreesurferMesh(fslmesh.Mesh):
@@ -81,8 +139,7 @@ class FreesurferMesh(fslmesh.Mesh):
             read_stamp=True)
 
         filename = op.abspath(filename)
-        name     = fslpath.removeExt(op.basename(filename),
-                                     GEOMETRY_EXTENSIONS)
+        name     = op.basename(filename)
 
         fslmesh.Mesh.__init__(self,
                               indices,
@@ -91,13 +148,15 @@ class FreesurferMesh(fslmesh.Mesh):
 
         self.addVertices(vertices, filename, fixWinding=fixWinding)
 
+        self.__luts = collections.OrderedDict()
+
         self.setMeta('comment', comment)
         for k, v in meta.items():
             self.setMeta(k, v)
 
         if loadAll:
 
-            allFiles = relatedFiles(filename, ftypes=GEOMETRY_EXTENSIONS)
+            allFiles = relatedGeometryFiles(filename)
 
             for f in allFiles:
                 verts, idxs = nibfs.read_geometry(f)
@@ -110,9 +169,9 @@ class FreesurferMesh(fslmesh.Mesh):
         ``nibabel.freesurfer.load_geometry``. Otherwise, it is passed to
         :meth:`.Mesh.loadVertices`.
         """
-        if not fslpath.hasExt(infile, GEOMETRY_EXTENSIONS):
-            return fslmesh.Mesh.loadVertices(
-                self, infile, key, **kwargs)
+
+        if not isGeometryFile(infile):
+            return fslmesh.Mesh.loadVertices(self, infile, key, **kwargs)
 
         infile = op.abspath(infile)
         if key is None:
@@ -133,68 +192,198 @@ class FreesurferMesh(fslmesh.Mesh):
     def loadVertexData(self, infile, key=None):
         """Overrides :meth:`.Mesh.loadVertexData`. If the given ``infile``
         looks like a Freesurfer file, it is loaded via the
-        :func:`loadFreesurferVertexFile` function. Otherwise, it is passed
-        through to the base-class function.
+        :func:`loadVertexDataFile` function. Otherwise, it is passed through
+        to the base-class function.
+
+        If the given ``infile`` is a vertex annotation file, it is assumed
+        that the file contains a value for every vertex in the mesh.
         """
 
-        if not fslpath.hasExt(infile, VERTEX_DATA_EXTENSIONS):
-            return fslmesh.loadVertexData(infile, key)
+        isvdata  = isVertexDataFile( infile)
+        isvlabel = isVertexLabelFile(infile)
+        isvannot = isVertexAnnotFile(infile)
 
-        infile = op.abspath(infile)
+        if not any((isvdata, isvlabel, isvannot)):
+            return fslmesh.Mesh.loadVertexData(infile)
+
+        infile    = op.abspath(infile)
+        nvertices = self.vertices.shape[0]
 
         if key is None:
             key = infile
 
-        vdata = loadFreesurferVertexFile(infile)
+        vdata = loadVertexDataFile(infile)
+
+        if isvlabel:
+            idxs, vdata    = np.asarray(vdata, np.int)
+            expanded       = np.zeros(nvertices)
+            expanded[idxs] = vdata
+            vdata          = expanded
+
+        elif isvannot:
+            vdata, lut, names = vdata
 
         self.addVertexData(key, vdata)
 
+        if isvlabel:
+            self.__luts[key] = lut, names, lut
+
         return vdata
 
 
-def loadFreesurferVertexFile(infile):
-    """Loads the given Freesurfer file, assumed to contain vertex-wise data.
+    def getVertexDataColourTable(self, key):
+        """TODO
+        """
+
+        return self.__luts[key]
+
+
+def loadVertexDataFile(infile):
+    """Loads the given Freesurfer vertex data, label, or annotation file.
+
+    This function return different things depending on what ``infile`` is:
+
+     - If ``infile`` is a vertex data file, a ``(nvertices,)`` array is
+       returned, containing one value for each vertex in the mesh.
+
+     - If ``infile`` is a ``mgh``/``mgz`` file, the image data is returned
+       as-is, with dimensions of length 1 squeezed out (under the assumption
+       that the image contains scalar vertex data).
+
+     - If ``infile`` is a vertex label file, a tuple containing the following
+       is returned:
+
+       - a ``(n,)`` array, containing the indices of all vertices that are
+         specified in the file.
+
+       - a ``(n,)`` array, containing scalar value for each vertex
+
+     - If ``infile`` is a vertex annotation file, a tuple containing the
+       following is returned:
+
+       - a ``(n,)`` array  containing the indices of all ``n`` vertices that
+         are specified in the file.
+
+       - a ``(l, 5)`` array containing the RGBA colour, and the label value,
+         for every label that is specified in the file.
+
+       - A list of length ``l``, containing the names of every label that is
+         specified in the file.
+
     """
 
-    # morphometry file
-    morphexts = ['.curv', '.crv', '.area', '.thickness', '.volume']
+    if isVertexDataFile(infile):
+        return nibfs.read_morph_data(infile)
 
-    if fslpath.hasExt(infile, morphexts):
-        vdata = nibfs.read_morph_data(infile)
+    elif isVertexLabelFile(infile):
+        return nibfs.read_label(infile, read_scalars=True)
+
+    elif isVertexAnnotFile(infile):
+        return nibfs.read_annot(infile, orig_ids=False)
 
-    # MGH image file
     elif fslpath.hasExt(infile, fslmgh.ALLOWED_EXTENSIONS):
-        vdata = fslmgh.MGHImage(infile)[:].squeeze()
+        return fslmgh.MGHImage(infile)[:].squeeze()
+    else:
+        raise ValueError('Unrecognised freesurfer '
+                         'file type: {}'.format(infile))
+
+
+@memoize.memoize
+def isCoreGeometryFile(infile):
+    """Returns ``True`` if ``infile`` looks like a core Freesurfer geometry
+    file, ``False`` otherwise.
+    """
+    infile = op.basename(infile)
+    return any([fnmatch.fnmatch(infile, gf) for gf in CORE_GEOMETRY_FILES])
+
+
+@memoize.memoize
+def isGeometryFile(infile):
+    """Returns ``True`` if ``infile`` looks like a Freesurfer geometry
+    file (core or otherwise), ``False`` otherwise.
+    """
+    infile = op.basename(infile)
+    return any([fnmatch.fnmatch(infile, gf)
+                for gf in CORE_GEOMETRY_FILES + EXTRA_GEOMETRY_FILES])
+
+
+@memoize.memoize
+def isVertexDataFile(infile):
+    """Returns ``True`` if ``infile`` looks like a Freesurfer vertex
+    data file, ``False`` otherwise.
+    """
+    infile = op.basename(infile)
+    return any([fnmatch.fnmatch(infile, gf) for gf in VERTEX_DATA_FILES])
+
+
+@memoize.memoize
+def isVertexLabelFile(infile):
+    """Returns ``True`` if ``infile`` looks like a Freesurfer vertex
+    label file, ``False`` otherwise.
+    """
+    infile = op.basename(infile)
+    return any([fnmatch.fnmatch(infile, gf) for gf in VERTEX_LABEL_FILES])
+
+
+@memoize.memoize
+def isVertexAnnotFile(infile):
+    """Returns ``True`` if ``infile`` looks like a Freesurfer vertex
+    annotation file, ``False`` otherwise.
+    """
+    infile = op.basename(infile)
+    return any([fnmatch.fnmatch(infile, gf) for gf in VERTEX_ANNOT_FILES])
+
+
+def relatedGeometryFiles(fname):
+    """Returns a list of all files which (look like they) are freesurfer
+    geometry files which correspond to the given geometry file.
+    """
+
+    if not isCoreGeometryFile(fname):
+        return []
+
+    dirname, fname = op.split(op.abspath(fname))
+    hemi           = fname[0]
+
+    fpats          = [hemi + p[1:] for p in CORE_GEOMETRY_FILES]
+
+    related        = [glob.glob(op.join(dirname, p)) for p in fpats]
+    related        = list(it.chain(*related))
+
+    return [r for r in related if op.basename(r) != fname]
 
-    return vdata
 
+def relatedVertexDataFiles(fname):
+    """Returns a list of all files which (look like they) are vertex data,
+    label, or annotation files related to the given freesurfer geometry file.
+    """
+
+    if not isCoreGeometryFile(fname):
+        return []
+
+    fname   = op.abspath(fname)
+    dirname = op.dirname(fname)
+    hemi    = op.basename(fname)[0]
+
+    fpats   = VERTEX_DATA_FILES + VERTEX_LABEL_FILES + VERTEX_ANNOT_FILES
+    fpats   = [hemi + p[1:] if p.startswith('?h') else p for p in fpats]
+
+    related = [glob.glob(op.join(dirname, p)) for p in fpats]
 
-def relatedFiles(fname, ftypes=None):
-    """Returns a list of all files which (look like they) are related to the
-    given freesurfer file.
+    return list(it.chain(*related))
+
+
+def findReferenceImage(fname):
+    """Attempts to locate the volumetric reference image for (what is
+    assumed to be) the given Freesurfer geometry file.
     """
 
-    if ftypes is None:
-        ftypes = VERTEX_DATA_EXTENSIONS
-
-    #
-    # .annot files contain labels for each vertex, and RGB values for each
-    #  label
-    #    -> nib.freesurfer.read_annot
-    #
-    # .label files contain scalar labels associated with each vertex
-    #    -> read_label
-    #
-    # .curv files contain vertex data
-    #    -> nib.freesurfer.read_morph_data
-    #
-    # .w files contain vertex data (potentially for a subset of vertices)
-    #    -> ?
-
-    prefix  = op.splitext(fname)[0]
-    related = []
-
-    for ftype in ftypes:
-        related += list(glob.glob('{}{}'.format(prefix, ftype)))
-
-    return [r for r in related if r != fname]
+    basedir = op.dirname(op.dirname(op.abspath(fname)))
+    t1      = op.join(basedir, 'mri', 'T1.mgz')
+    exts    = fslimage.ALLOWED_EXTENSIONS + fslmgh.ALLOWED_EXTENSIONS
+
+    try:
+
+        fslpath.addExt(t1, allowedExts=exts, mustExist=True)
+    except fslpath.PathError:
+        return None