forked from dipy/dipy
-
Notifications
You must be signed in to change notification settings - Fork 5
/
vtk.py
93 lines (75 loc) · 2.79 KB
/
vtk.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
from __future__ import division, print_function, absolute_import
from dipy.viz.utils import set_input
# Conditional import machinery for vtk
from dipy.utils.optpkg import optional_package
# Allow import, but disable doctests if we don't have vtk
vtk, have_vtk, setup_module = optional_package('vtk')
colors, have_vtk_colors, _ = optional_package('vtk.util.colors')
ns, have_numpy_support, _ = optional_package('vtk.util.numpy_support')
if have_vtk:
version = vtk.vtkVersion.GetVTKSourceVersion().split(' ')[-1]
major_version = vtk.vtkVersion.GetVTKMajorVersion()
def load_polydata(file_name):
""" Load a vtk polydata to a supported format file
Supported file formats are OBJ, VTK, FIB, PLY, STL and XML
Parameters
----------
file_name : string
Returns
-------
output : vtkPolyData
"""
# get file extension (type) lower case
file_extension = file_name.split(".")[-1].lower()
if file_extension == "vtk":
reader = vtk.vtkPolyDataReader()
elif file_extension == "fib":
reader = vtk.vtkPolyDataReader()
elif file_extension == "ply":
reader = vtk.vtkPLYReader()
elif file_extension == "stl":
reader = vtk.vtkSTLReader()
elif file_extension == "xml":
reader = vtk.vtkXMLPolyDataReader()
elif file_extension == "obj":
try: # try to read as a normal obj
reader = vtk.vtkOBJReader()
except: # than try load a MNI obj format
reader = vtk.vtkMNIObjectReader()
else:
raise "polydata " + file_extension + " is not suported"
reader.SetFileName(file_name)
reader.Update()
# print(file_name + " Mesh " + file_extension + " Loaded")
return reader.GetOutput()
def save_polydata(polydata, file_name, binary=False, color_array_name=None):
""" Save a vtk polydata to a supported format file
Save formats can be VTK, FIB, PLY, STL and XML.
Parameters
----------
polydata : vtkPolyData
file_name : string
"""
# get file extension (type)
file_extension = file_name.split(".")[-1].lower()
if file_extension == "vtk":
writer = vtk.vtkPolyDataWriter()
elif file_extension == "fib":
writer = vtk.vtkPolyDataWriter()
elif file_extension == "ply":
writer = vtk.vtkPLYWriter()
elif file_extension == "stl":
writer = vtk.vtkSTLWriter()
elif file_extension == "xml":
writer = vtk.vtkXMLPolyDataWriter()
elif file_extension == "obj":
raise "mni obj or Wavefront obj ?"
# writer = set_input(vtk.vtkMNIObjectWriter(), polydata)
writer.SetFileName(file_name)
writer = set_input(writer, polydata)
if color_array_name is not None:
writer.SetArrayName(color_array_name)
if binary:
writer.SetFileTypeToBinary()
writer.Update()
writer.Write()