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

Read our documentation on viewing source code .

Loading