Ticket #5116: VTKBlender.py

File VTKBlender.py, 11.2 KB (added by anonymous, 8 years ago)
Line 
1# $Id: VTKBlender.py,v 1.19 2008-07-03 15:13:21 cwant Exp $
2#
3# Copyright (c) 2005, Chris Want, Research Support Group,
4# AICT, University of Alberta. All rights reserved.
5#
6# Redistribution and use in source and binary forms, with or without
7# modification, are permitted provided that the following conditions are met:
8#
9# 1) Redistributions of source code must retain the above copyright
10# notice, this list of conditions and the following disclaimer.
11# 2) Redistributions in binary form must reproduce the above copyright
12# notice, this list of conditions and the following disclaimer in the
13# documentation and/or other materials provided with the distribution.
14#
15# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
16# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
17# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
18# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE
19# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
20# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
21# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
22# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
23# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
24# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF
25# THE POSSIBILITY OF SUCH DAMAGE.
26#
27# Contributors: Chris Want (University of Alberta),
28# Fritz Mielert (University of Stuttgart)
29"""
30VTK inside Blender module.
31
32 This module provides code so that polydata from vtk can
33 be used inside of blender.
34
35 Python needs to find the vtk stuff and this module in order
36 for this to work, and you can either a) set your PYTHONPATH
37 in your environment, or you can b) hardcode your vtk path's
38 in your script, e.g.,
39
40 a) at the prompt, before starting blender, type:
41 PYTHONPATH=$VTK_ROOT/Wrapping/Python:${LIBRARY_OUTPUT_PATH}
42 PYTHONPATH=$PYTHONPATH:${PATH_TO_THIS_MODULE}
43 export PYTHONPATH
44
45 b) add the following to your script near the beginning, before
46 importing vtk or VTKBlender:
47
48 import sys
49 sys.path.append($VTK_ROOT/Wrapping/Python)
50 sys.path.append(${LIBRARY_OUTPUT_PATH})
51 sys.path.append(${PATH_TO_VTKBlender_MODULE})
52
53 Be sure to replace $VTK_ROOT and ${LIBRARY_OUTPUT_PATH} with
54 values that are relevant to your system. These values can be
55 found by starting vtkpython with no arguments and typing:
56
57 import sys
58 print sys.path
59
60 Usually the first two items reported are the ones you want.
61
62 Also replace ${PATH_TO_VTKBlender_MODULE} with wherever you have
63 put the VTKBlender module.
64"""
65
66import vtk
67import time, string
68
69try:
70 import Blender
71 from Blender import Mesh, Object, Material
72
73except:
74 print "No Blender module found!"
75
76__versiontag__ = "$Revision: 1.19 $"
77__version__ = string.split(__versiontag__)[1]
78
79# some flags to alter behavior
80flags = 0
81TRIS_TO_QUADS = 0x01
82SMOOTH_FACES = 0x02
83
84# What is this 'tri to quad' stuff? Well, sometimes it's best to
85# try to read in pairs of consecutive triangles in as quad faces.
86# An example: you extrude a tube along a polyline in vtk, and if
87# you can get it into Blender as a bunch of quads, you can use a
88# Catmull-Clark subdivision surface to smooth the tube out, with
89# fewer creases.
90
91def SetTrisToQuads():
92 global flags
93 flags = flags | TRIS_TO_QUADS
94
95def SetTrisToTris():
96 global flags
97 flags = flags & ~TRIS_TO_QUADS
98
99def SetFacesToSmooth():
100 global flags
101 flags = flags | SMOOTH_FACES
102
103def SetFacesToFaceted():
104 global flags
105 flags = flags & ~SMOOTH_FACES
106
107def BlenderToPolyData(me, uvlayer=None):
108
109 pcoords = vtk.vtkFloatArray()
110 pcoords.SetNumberOfComponents(3)
111 pcoords.SetNumberOfTuples(len(me.verts))
112 for i in range(len(me.verts)):
113 p0 = me.verts[i].co[0]
114 p1 = me.verts[i].co[1]
115 p2 = me.verts[i].co[2]
116 pcoords.SetTuple3(i, p0, p1, p2)
117
118 points = vtk.vtkPoints()
119 points.SetData(pcoords)
120
121 polys = vtk.vtkCellArray()
122 lines = vtk.vtkCellArray()
123 for face in me.faces:
124 if len(face.v) == 4:
125 polys.InsertNextCell(4)
126 polys.InsertCellPoint(face.v[0].index)
127 polys.InsertCellPoint(face.v[1].index)
128 polys.InsertCellPoint(face.v[2].index)
129 polys.InsertCellPoint(face.v[3].index)
130 elif len(face.v) == 3:
131 polys.InsertNextCell(3)
132 polys.InsertCellPoint(face.v[0].index)
133 polys.InsertCellPoint(face.v[1].index)
134 polys.InsertCellPoint(face.v[2].index)
135 elif len(face.v) == 2:
136 lines.InsertNextCell(2)
137 lines.InsertCellPoint(face.v[0].index)
138 lines.InsertCellPoint(face.v[1].index)
139
140 for edge in me.edges:
141 lines.InsertNextCell(2)
142 lines.InsertCellPoint(edge.v1.index)
143 lines.InsertCellPoint(edge.v2.index)
144
145 pdata =vtk.vtkPolyData()
146 pdata.SetPoints(points)
147 pdata.SetPolys(polys)
148 pdata.SetLines(lines)
149
150 if me.faceUV:
151 if uvlayer:
152 uvnames = me.getUVLayerNames()
153 if uvlayer in uvnames:
154 me.activeUVLayer = uvlayer
155 tcoords = vtk.vtkFloatArray()
156 tcoords.SetNumberOfComponents(2)
157 tcoords.SetNumberOfTuples(len(me.verts))
158 for face in me.faces:
159 for i in range(len(face.verts)):
160 uv = face.uv[i]
161 tcoords.SetTuple2(face.v[i].index, uv[0], uv[1])
162 pdata.GetPointData().SetTCoords(tcoords);
163
164 pdata.Update()
165
166 return pdata
167
168def PolyDataMapperToBlender(pmapper, me=None):
169 global flags
170 faces = []
171 edges = []
172 oldmats = None
173
174 newmesh = 0
175 if (me == None):
176 me = Mesh.New()
177 newmesh = 1
178 else:
179 if me.materials:
180 oldmats = me.materials
181
182 me.verts = None # this kills the faces/edges tooo
183
184 pmapper.Update()
185
186 pdata = pmapper.GetInput()
187 plut = pmapper.GetLookupTable()
188 #print pdata.GetNumberOfCells()
189
190 scalars = pdata.GetPointData().GetScalars()
191
192 verts = []
193 for i in range(pdata.GetNumberOfPoints()):
194 point = pdata.GetPoint(i)
195 verts.append([point[0],point[1],point[2]])
196
197 me.verts.extend(verts)
198 # I think we can free some memory by killing the reference
199 # from vert to the list it points at (not sure though)
200 verts = []
201
202 colors = None
203
204 if ( (scalars != None) and (plut != None) ):
205 colors = []
206
207 # Have to be a bit careful since VTK 5.0 changed the
208 # prototype of vtkLookupTable.GetColor()
209 try:
210 # VTK 5.x
211 scolor = [0,0,0]
212 for i in range(scalars.GetNumberOfTuples()):
213 plut.GetColor(scalars.GetTuple1(i), scolor)
214 color = map(VTKToBlenderColor, scolor)
215 alpha = int(plut.GetOpacity(scalars.GetTuple1(i))*255)
216 colors.append([color[0], color[1], color[2], alpha])
217
218 except:
219 # VTK 4.x
220 for i in range(scalars.GetNumberOfTuples()):
221 color = map(VTKToBlenderColor, \
222 plut.GetColor(scalars.GetTuple1(i)))
223 alpha = int(plut.GetOpacity(scalars.GetTuple1(i))*255)
224 colors.append([color[0], color[1], color[2], alpha])
225
226 skiptriangle = False
227 for i in range(pdata.GetNumberOfCells()):
228
229 cell = pdata.GetCell(i)
230
231 #print i, pdata.GetCellType(i)
232
233 # Do lines
234 if pdata.GetCellType(i)==3:
235
236 n1 = cell.GetPointId(0)
237 n2 = cell.GetPointId(1)
238
239 BlenderAddEdge(me, edges, n1, n2)
240
241 # Do poly lines
242 if pdata.GetCellType(i)==4:
243 for j in range(cell.GetNumberOfPoints()-1):
244
245 n1 = cell.GetPointId(j)
246 n2 = cell.GetPointId(j+1)
247
248 BlenderAddEdge(me, edges, n1, n2)
249
250 # Do triangles
251 if pdata.GetCellType(i)==5:
252 if skiptriangle==True:
253 skiptriangle = False
254 elif ( (flags & TRIS_TO_QUADS) and
255 (i < pdata.GetNumberOfCells()-1) and
256 (pdata.GetCellType(i+1)==5) ):
257 n1 = cell.GetPointId(0)
258 n2 = cell.GetPointId(1)
259 n3 = cell.GetPointId(2)
260 nextcell = pdata.GetCell(i+1)
261 m1 = nextcell.GetPointId(0)
262 m2 = nextcell.GetPointId(1)
263 m3 = nextcell.GetPointId(2)
264
265 if ( (n2 == m3) and (n3 == m2) ):
266 BlenderAddFace(me, faces, n1, n2, m1, n3)
267 skiptriangle = True
268 else:
269 BlenderAddFace(me, faces, n1, n2, n3)
270
271 else:
272 n1 = cell.GetPointId(0)
273 n2 = cell.GetPointId(1)
274 n3 = cell.GetPointId(2)
275
276 BlenderAddFace(me, faces, n1, n2, n3)
277
278 # Do triangle strips
279 if pdata.GetCellType(i)==6:
280 numpoints = cell.GetNumberOfPoints()
281 if ( (flags & TRIS_TO_QUADS) and (numpoints % 2 == 0) ):
282 for j in range(cell.GetNumberOfPoints()-3):
283 if (j % 2 == 0):
284 n1 = cell.GetPointId(j)
285 n2 = cell.GetPointId(j+1)
286 n3 = cell.GetPointId(j+2)
287 n4 = cell.GetPointId(j+3)
288
289 BlenderAddFace(me, faces, n1, n2, n4, n3)
290 else:
291 for j in range(cell.GetNumberOfPoints()-2):
292 if (j % 2 == 0):
293 n1 = cell.GetPointId(j)
294 n2 = cell.GetPointId(j+1)
295 n3 = cell.GetPointId(j+2)
296 else:
297 n1 = cell.GetPointId(j)
298 n2 = cell.GetPointId(j+2)
299 n3 = cell.GetPointId(j+1)
300
301 BlenderAddFace(me, faces, n1, n2, n3)
302 # Do polygon
303 if pdata.GetCellType(i)==7:
304 # Add a vert at the center of the polygon,
305 # and break into triangles
306 x = 0.0
307 y = 0.0
308 z = 0.0
309 scal = 0.0
310 N = cell.GetNumberOfPoints()
311 for j in range(N):
312 point = pdata.GetPoint(cell.GetPointId(j))
313 x = x + point[0]
314 y = y + point[1]
315 z = z + point[2]
316 if (scalars != None):
317 scal = scal + scalars.GetTuple1(j)
318 x = x / N
319 y = y / N
320 z = z / N
321 scal = scal / N
322
323 newidx = len(me.verts)
324 me.verts.extend(x,y,z)
325
326 if (scalars != None):
327 try:
328 # VTK 5.x
329 scolor = [0,0,0]
330 plut.GetColor(scal, scolor)
331 color = map(VTKToBlenderColor, scolor)
332 except:
333 color = map(VTKToBlenderColor, plut.GetColor(scal))
334 alpha = int(plut.GetOpacity(scalars.GetTuple1(i))*255)
335
336 colors.append([color[0], color[1], color[2], alpha])
337
338 # Add triangles connecting polynomial sides to new vert
339 for j in range(N):
340 n1 = cell.GetPointId(j)
341 n2 = cell.GetPointId( (j+1) % N )
342 n3 = newidx
343 BlenderAddFace(me, faces, n1, n2, n3)
344
345 # Do pixel
346 if pdata.GetCellType(i)==8:
347 n1 = cell.GetPointId(0)
348 n2 = cell.GetPointId(1)
349 n3 = cell.GetPointId(2)
350 n4 = cell.GetPointId(3)
351
352 BlenderAddFace(me, faces, n1, n2, n3, n4)
353 # Do quad
354 if pdata.GetCellType(i)==9:
355 n1 = cell.GetPointId(0)
356 n2 = cell.GetPointId(1)
357 n3 = cell.GetPointId(2)
358 n4 = cell.GetPointId(3)
359
360 BlenderAddFace(me, faces, n1, n2, n3, n4)
361
362 if len(edges) > 0:
363 me.edges.extend(edges)
364 if len(faces) > 0:
365 me.faces.extend(faces)
366
367 if ( flags & SMOOTH_FACES):
368 for f in me.faces:
369 f.smooth = 1
370
371 # Some faces in me.faces may have been discarded from our
372 # list, so best to compute the vertex colors after the faces
373 # have been added to the mesh
374 if (colors != None):
375 me.vertexColors = 1
376 for f in me.faces:
377 f_col = []
378 for v in f.v:
379 f_col.append(colors[v.index])
380
381 SetVColors(f.col, f_col)
382
383 if not me.materials:
384 if oldmats:
385 me.materials = oldmats
386 else:
387 newmat = Material.New()
388 if (colors != None):
389 newmat.mode |= Material.Modes.VCOL_PAINT
390 me.materials = [newmat]
391
392 if (newmesh==0):
393 me.update()
394
395 return me
396
397def VTKToBlenderColor(x):
398 return int(255*float(x)+0.5)
399
400def BlenderAddFace(me, faces, n1, n2, n3, n4=None):
401
402 if (n4 != None):
403 faces.append([me.verts[n1], me.verts[n2], \
404 me.verts[n3], me.verts[n4]])
405 else:
406 faces.append([me.verts[n1], me.verts[n2], me.verts[n3]])
407
408def BlenderAddEdge(me, edges, n1, n2):
409 edges.append([me.verts[n1], me.verts[n2]])
410
411
412def SetVColors(col, vcols):
413 for j in range(len(col)):
414
415 col[j].r = vcols[j][0]
416 col[j].g = vcols[j][1]
417 col[j].b = vcols[j][2]
418 if len(vcols[j]) == 3:
419 col[j].a = 255
420 else:
421 col[j].a = vcols[j][3]
422