1
|
16
|
__all__ = [
|
2
|
|
'check_orientation',
|
3
|
|
'check_orthogonal',
|
4
|
|
'add_data',
|
5
|
|
'add_textures',
|
6
|
|
]
|
7
|
|
|
8
|
|
|
9
|
16
|
import numpy as np
|
10
|
16
|
import vtk
|
11
|
16
|
from vtk.util import numpy_support as nps
|
12
|
16
|
import pyvista
|
13
|
16
|
from PIL import Image
|
14
|
|
|
15
|
16
|
try:
|
16
|
16
|
from pyvista import is_pyvista_obj as is_pyvista_dataset
|
17
|
16
|
except ImportError:
|
18
|
16
|
from pyvista import is_pyvista_dataset
|
19
|
|
|
20
|
|
|
21
|
16
|
def check_orientation(axis_u, axis_v, axis_w):
|
22
|
|
"""This will check if the given ``axis_*`` vectors are the typical
|
23
|
|
cartesian refernece frame (i.e. rectilinear).
|
24
|
|
"""
|
25
|
16
|
if ( np.allclose(axis_u, (1, 0, 0)) and
|
26
|
|
np.allclose(axis_v, (0, 1, 0)) and
|
27
|
|
np.allclose(axis_w, (0, 0, 1)) ):
|
28
|
16
|
return True
|
29
|
16
|
return False
|
30
|
|
|
31
|
|
|
32
|
16
|
def check_orthogonal(axis_u, axis_v, axis_w):
|
33
|
|
"""Makes sure that the three input vectors are orthogonal"""
|
34
|
14
|
if not (np.abs(axis_u.dot(axis_v) < 1e-6) and
|
35
|
|
np.abs(axis_v.dot(axis_w) < 1e-6) and
|
36
|
|
np.abs(axis_w.dot(axis_u) < 1e-6)):
|
37
|
|
#raise ValueError('axis_u, axis_v, and axis_w must be orthogonal')
|
38
|
0
|
return False
|
39
|
14
|
return True
|
40
|
|
|
41
|
|
|
42
|
16
|
def add_data(output, data):
|
43
|
|
"""Adds data arrays to an output VTK data object"""
|
44
|
16
|
for d in data:
|
45
|
16
|
arr = d.array.array
|
46
|
16
|
c = nps.numpy_to_vtk(num_array=arr)
|
47
|
16
|
c.SetName(d.name)
|
48
|
16
|
loc = d.location
|
49
|
16
|
if loc == 'vertices':
|
50
|
16
|
output.GetPointData().AddArray(c)
|
51
|
|
else:
|
52
|
16
|
output.GetCellData().AddArray(c)
|
53
|
16
|
return output
|
54
|
|
|
55
|
|
|
56
|
16
|
def add_textures(output, textures, elname):
|
57
|
|
"""Add textures to a pyvista data object"""
|
58
|
16
|
if not is_pyvista_dataset(output):
|
59
|
16
|
output = pyvista.wrap(output)
|
60
|
|
|
61
|
14
|
for i, tex in enumerate(textures):
|
62
|
|
# Now map the coordinates for the texture
|
63
|
14
|
m = vtk.vtkTextureMapToPlane()
|
64
|
14
|
m.SetInputDataObject(output)
|
65
|
14
|
m.SetOrigin(tex.origin)
|
66
|
14
|
m.SetPoint1(tex.origin + tex.axis_u)
|
67
|
14
|
m.SetPoint2(tex.origin + tex.axis_v)
|
68
|
14
|
m.Update()
|
69
|
|
# Grab the texture coordinates
|
70
|
14
|
tmp = m.GetOutputDataObject(0)
|
71
|
14
|
tcoord = tmp.GetPointData().GetTCoords()
|
72
|
14
|
name = tex.name
|
73
|
14
|
if name is None or name == '':
|
74
|
14
|
name = '{}-texture-{}'.format(elname, i)
|
75
|
14
|
tcoord.SetName(name)
|
76
|
|
# Add these coordinates to the PointData of the output
|
77
|
|
# NOTE: Let pyvista handle setting the TCoords because of how VTK cleans
|
78
|
|
# up old TCoords
|
79
|
14
|
output.GetPointData().AddArray(tcoord)
|
80
|
|
# Add the vtkTexture to the output
|
81
|
14
|
img = np.array(Image.open(tex.image))
|
82
|
14
|
tex.image.seek(0) # Reset the image bytes in case it is accessed again
|
83
|
14
|
if img.shape[2] > 3:
|
84
|
0
|
img = img[:, :, 0:3]
|
85
|
14
|
vtexture = pyvista.numpy_to_texture(img)
|
86
|
14
|
output.textures[name] = vtexture
|
87
|
14
|
output._activate_texture(name)
|
88
|
|
|
89
|
14
|
return output
|