Source code for fsl.data.vtk

#
# 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.

A handful of convenience functions are also in this module:

.. autosummary::
   :nosignatures:

   loadVTKPolydataFile
   getFIRSTPrefix
   findReferenceImage


.. 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`."""


[docs]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=name, dataSource=dataSource, vertices=data, fixWinding=fixWinding)
[docs]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
[docs]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
[docs]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