mesh.py 8.72 KB
Newer Older
1
2
#!/usr/bin/env python
#
3
# mesh.py - The TriangleMesh class.
4
5
6
#
# Author: Paul McCarthy <pauldmccarthy@gmail.com>
#
7
8
"""This module provides the :class:`TriangleMesh` class, which represents a
3D model made of triangles.
9

10
.. note:: I/O support is very limited - currently, the only supported file
11
          type is the VTK legacy file format, containing the ``POLYDATA``
12
13
14
          dataset. the :class:`TriangleMesh` class assumes that every polygon
          defined in an input file is a triangle (i.e. refers to three
          vertices).
15
16
17

          See http://www.vtk.org/wp-content/uploads/2015/04/file-formats.pdf
          for an overview of the VTK legacy file format.
18
19

          In the future, I may or may not add support for more complex meshes.
20
21
"""

22

23
24
import logging

25
26
import os.path as op
import numpy   as np
27

Paul McCarthy's avatar
Paul McCarthy committed
28
29
import six

30
31
from . import image as fslimage

32
33
34
35

log = logging.getLogger(__name__)


36
37
38
class TriangleMesh(object):
    """The ``TriangleMesh`` class represents a 3D model. A mesh is defined by
    a collection of vertices and indices.  The indices index into the list of
39
    vertices, and define a set of triangles which make the model.
40
41
42
43


    A ``TriangleMesh`` instance has the following attributes:

44

45
46
    ============== ====================================================
    ``name``       A name, typically the file name sans-suffix.
47

48
49
    ``dataSource`` Full path to the mesh file (or ``None`` if there is
                   no file associated with this mesh).
50

Paul McCarthy's avatar
Paul McCarthy committed
51
    ``vertices``   A :math:`N\times 3` ``numpy`` array containing
52
                   the vertices.
53

54
55
56
57
    ``indices``    A :meth:`M\times 3` ``numpy`` array containing
                   the vertex indices for :math:`M` triangles
    ============== ====================================================

58
59
60
61
62
63
64
65

    And the following methods:

    .. autosummary::
       :nosignatures:

       getBounds
       loadVertexData
66
67
       getVertexData
       clearVertexData
68
    """
69

70

71
    def __init__(self, data, indices=None):
72
        """Create a ``TriangleMesh`` instance.
73

74
75
76
77
        :arg data:    Can either be a file name, or a :math:`N\\times 3`
                      ``numpy`` array containing vertex data. If ``data`` is
                      a file name, it is passed to the
                      :func:`loadVTKPolydataFile` function.
78

79
80
        :arg indices: A list of indices into the vertex data, defining
                      the triangles.
81
82
        """

Paul McCarthy's avatar
Paul McCarthy committed
83
        if isinstance(data, six.string_types):
84
85
86
87
88
89
            infile = data
            data, lengths, indices = loadVTKPolydataFile(infile)

            if np.any(lengths != 3):
                raise RuntimeError('All polygons in VTK file must be '
                                   'triangles ({})'.format(infile))
90
91
92
93

            self.name       = op.basename(infile)
            self.dataSource = infile
        else:
94
95
            self.name       = 'TriangleMesh'
            self.dataSource = None
96

97
        if indices is None:
98
            indices = np.arange(data.shape[0])
99

100
101
        self.vertices     = np.array(data)
        self.indices      = np.array(indices).reshape((-1, 3))
102

103
104
105
        self.__vertexData = {}
        self.__loBounds   = self.vertices.min(axis=0)
        self.__hiBounds   = self.vertices.max(axis=0)
106

107

108
    def __repr__(self):
109
110
        """Returns a string representation of this ``TriangleMesh`` instance.
        """
111
112
113
114
115
        return '{}({}, {})'.format(type(self).__name__,
                                   self.name,
                                   self.dataSource)

    def __str__(self):
116
117
        """Returns a string representation of this ``TriangleMesh`` instance.
        """
118
119
120
        return self.__repr__()


121
    def getBounds(self):
122
        """Returns a tuple of values which define a minimal bounding box that
123
        will contain all vertices in this ``TriangleMesh`` instance. The
124
        bounding box is arranged like so:
125
126
127

            ``((xlow, ylow, zlow), (xhigh, yhigh, zhigh))``
        """
128
        return (self.__loBounds, self.__hiBounds)
129
130


131
    def loadVertexData(self, dataSource, vertexData=None):
132
        """Attempts to load scalar data associated with each vertex of this
133
134
135
        ``TriangleMesh`` from the given ``dataSource``. The data is returned,
        and also stored in an internal cache so it can be retrieved later
        via the :meth:`getVertexData` method.
136
137

        This method may be overridden by sub-classes.
138
139
140
141

        :arg dataSource: Path to the vertex data to load
        :arg vertexData: The vertex data itself, if it has already been
                         loaded.
142
143
144

        :returns: A ``(M, N)``) array, which contains ``N`` data points
                  for ``M`` vertices.
145
        """
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171

        nvertices = self.vertices.shape[0]

        # Currently only white-space delimited
        # text files are supported
        if vertexData is None:
            vertexData = np.loadtxt(dataSource)
            vertexData.reshape(nvertices, -1)

        if vertexData.shape[0] != nvertices:
            raise ValueError('Incompatible size: {}'.format(dataSource))

        self.__vertexData[dataSource] = vertexData

        return vertexData


    def getVertexData(self, dataSource):
        """Returns the vertex data for the given ``dataSource`` from the
        internal vertex data cache. If the given ``dataSource`` is not
        in the cache, it is loaded via :meth:`loadVertexData`.
        """

        try:             return self.__vertexData[dataSource]
        except KeyError: return self.loadVertexData(dataSource)

172

173
174
175
176
177
178
    def clearVertexData(self):
        """Clears the internal vertex data cache - see the
        :meth:`loadVertexData` and :meth:`getVertexData`  methods.
        """

        self.__vertexData = {}
179
180


181
ALLOWED_EXTENSIONS     = ['.vtk']
182
183
"""A list of file extensions which could contain :class:`TriangleMesh` data.
"""
184
185
186
187
188
189
190
191
192
193
194
195


EXTENSION_DESCRIPTIONS = ['VTK polygon model file']
"""A description for each of the extensions in :data:`ALLOWED_EXTENSIONS`."""


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:
196

197
198
199
200
201
202
                - 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.
    """
203

204
205
206
207
208
209
210
211
212
    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')
213

214
215
    nVertices = int(lines[4].split()[1])
    nPolygons = int(lines[5 + nVertices].split()[1])
216
217
    nIndices  = int(lines[5 + nVertices].split()[2]) - nPolygons

218
219
220
221
222
223
224
    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, :] = map(float, vertLine.split())
Paul McCarthy's avatar
Paul McCarthy committed
225
        vertices[i, :] = [float(w) for w in vertLine.split()]
226
227
228
229
230
231
232
233
234
235

    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] = map(int, polyLine[1:])
Paul McCarthy's avatar
Paul McCarthy committed
236
        indices[start:end] = [int(w) for w in polyLine[1:]]
237
238
239
240

        indexOffset        += polygonLengths[i]

    return vertices, polygonLengths, indices
241
242
243
244
245


def getFIRSTPrefix(modelfile):
    """If the given ``vtk`` file was generated by `FIRST
    <https://fsl.fmrib.ox.ac.uk/fsl/fslwiki/FIRST>`_, this function
246
    will return the file prefix. Otherwise a ``ValueError`` will be
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
    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:

271
272
        dirname  = op.dirname(modelfile)
        prefixes = [getFIRSTPrefix(modelfile)]
273
274
    except:
        return None
275
276
277
278
279
280
281
282
283
284
285

    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:
            continue

    return None