VTK stuff moved to its own module. Also initial (non-working) freesurfer module

fsl/data/freesurfer.py

+#!/usr/bin/env python
+#
+# freesurfer.py -
+#
+# Author: Paul McCarthy <pauldmccarthy@gmail.com>
+#
+"""
+"""
+
+
+import os.path as op
+
+import nibabel as nib
+
+import fsl.utils.path as fslpath
+import fsl.data.mesh  as fslmesh
+
+
+class FreesurferMesh(fslmesh.TriangleMesh):
+    """
+    """
+
+    def __init__(self, filename):
+        """
+        """
+
+        vertices, indices, meta, comment = nib.freesurfer.read_geometry(
+            filename,
+            read_metadata=True,
+            read_stamp=True)
+
+        fslmesh.TriangleMesh.__init__(self, vertices, indices)
+
+        self.dataSource = op.abspath(filename)
+        self.name       = fslpath.removeExt(op.basename(filename),
+                                            ALLOWED_EXTENSIONS)
+
+
+
+    def loadVertexData(self, dataSource, vertexData=None):
+        pass
+
+
+
+ALLOWED_EXTENSIONS = ['.pial',
+                      '.white',
+                       '.sphere',
+                      '.inflated',
+                      '.orig',
+                      '.sulc',
+                      '.mid' ]
+EXTENSION_DESCRIPTIONS = [
+]
+
+
+def relatedFiles(fname):
+    """
+    """
+
+    #
+    # .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)
+    #    -> ?
+    pass

fsl/data/vtk.py

+#!/usr/bin/env python
+#
+# vtk.py - The VTKMesh class.
+#
+# Author: Paul McCarthy <pauldmccarthy@gmail.com>
+#
+"""This module provides the :class:`VTKMesh` class, for loading triangle
+meshes from VTK files.
+
+
+.. note:: I/O support is very limited - currently, the only supported file
+          type is the VTK legacy file format, containing the ``POLYDATA``
+          dataset. the :class:`TriangleMesh` class assumes that every polygon
+          defined in an input file is a triangle (i.e. refers to three
+          vertices).
+
+          See http://www.vtk.org/wp-content/uploads/2015/04/file-formats.pdf
+          for an overview of the VTK legacy file format.
+
+          In the future, I may or may not add support for more complex meshes.
+"""
+
+
+import os.path as op
+
+import numpy as np
+
+import fsl.data.mesh  as fslmesh
+import fsl.data.image as fslimage
+
+
+ALLOWED_EXTENSIONS     = ['.vtk']
+"""A list of file extensions which could contain :class:`VTKMesh` data.
+"""
+
+
+EXTENSION_DESCRIPTIONS = ['VTK polygon model file']
+"""A description for each of the extensions in :data:`ALLOWED_EXTENSIONS`."""
+
+
+class VTKMesh(fslmesh.Mesh):
+    """The ``VTKMesh`` class represents a triangle mesh loaded from a VTK
+    file. Typically only one set of vertices will be associated with a
+    ``VTKMesh``.
+    """
+
+
+    def __init__(self, infile, fixWinding=False):
+        """Create a ``VTKMesh``.
+
+        :arg infile:     VTK file to load mesh from.
+        :arg fixWinding: See the :meth:`.Mesh.addVertices` method.
+        """
+
+        data, lengths, indices = loadVTKPolydataFile(infile)
+        name                   = op.basename(infile)
+        dataSource             = op.abspath(infile)
+
+        if np.any(lengths != 3):
+            raise RuntimeError('All polygons in VTK file must be '
+                               'triangles ({})'.format(infile))
+
+        fslmesh.Mesh.__init__(self, indices, name, dataSource)
+
+        self.addVertices(data, 'default', fixWinding=fixWinding)
+
+
+def loadVTKPolydataFile(infile):
+    """Loads a vtk legacy file containing a ``POLYDATA`` data set.
+
+    :arg infile: Name of a file to load from.
+
+    :returns: a tuple containing three values:
+
+                - A :math:`N\\times 3` ``numpy`` array containing :math:`N`
+                  vertices.
+                - A 1D ``numpy`` array containing the lengths of each polygon.
+                - A 1D ``numpy`` array containing the vertex indices for all
+                  polygons.
+    """
+
+    lines = None
+
+    with open(infile, 'rt') as f:
+        lines = f.readlines()
+
+    lines = [l.strip() for l in lines]
+
+    if lines[3] != 'DATASET POLYDATA':
+        raise ValueError('Only the POLYDATA data type is supported')
+
+    nVertices = int(lines[4].split()[1])
+    nPolygons = int(lines[5 + nVertices].split()[1])
+    nIndices  = int(lines[5 + nVertices].split()[2]) - nPolygons
+
+    vertices       = np.zeros((nVertices, 3), dtype=np.float32)
+    polygonLengths = np.zeros( nPolygons,     dtype=np.uint32)
+    indices        = np.zeros( nIndices,      dtype=np.uint32)
+
+    for i in range(nVertices):
+        vertLine       = lines[i + 5]
+        vertices[i, :] = [float(w) for w in vertLine.split()]
+
+    indexOffset = 0
+    for i in range(nPolygons):
+
+        polyLine          = lines[6 + nVertices + i].split()
+        polygonLengths[i] = int(polyLine[0])
+
+        start              = indexOffset
+        end                = indexOffset + polygonLengths[i]
+        indices[start:end] = [int(w) for w in polyLine[1:]]
+
+        indexOffset        += polygonLengths[i]
+
+    return vertices, polygonLengths, indices
+
+
+def getFIRSTPrefix(modelfile):
+    """If the given ``vtk`` file was generated by `FIRST
+    <https://fsl.fmrib.ox.ac.uk/fsl/fslwiki/FIRST>`_, this function
+    will return the file prefix. Otherwise a ``ValueError`` will be
+    raised.
+    """
+
+    if not modelfile.endswith('first.vtk'):
+        raise ValueError('Not a first vtk file: {}'.format(modelfile))
+
+    modelfile = op.basename(modelfile)
+    prefix    = modelfile.split('-')
+    prefix    = '-'.join(prefix[:-1])
+
+    return prefix
+
+
+def findReferenceImage(modelfile):
+    """Given a ``vtk`` file, attempts to find a corresponding ``NIFTI``
+    image file. Return the path to the image, or ``None`` if no image was
+    found.
+
+    Currently this function will only return an image for ``vtk`` files
+    generated by `FIRST <https://fsl.fmrib.ox.ac.uk/fsl/fslwiki/FIRST>`_.
+    """
+
+    try:
+
+        dirname  = op.dirname(modelfile)
+        prefixes = [getFIRSTPrefix(modelfile)]
+    except ValueError:
+        return None
+
+    if prefixes[0].endswith('_first'):
+        prefixes.append(prefixes[0][:-6])
+
+    for p in prefixes:
+        try:
+            return fslimage.addExt(op.join(dirname, p), mustExist=True)
+        except fslimage.PathError:
+            continue
+
+    return None