Hot-keys on this page

r m x p   toggle line displays

j k   next/prev highlighted chunk

0   (zero) top of page

1   (one) first highlighted chunk

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

94

95

96

97

98

99

100

101

102

103

104

105

106

107

108

109

110

111

112

113

114

115

116

117

118

119

120

121

122

123

124

125

126

127

128

129

130

131

132

133

134

135

136

137

138

139

140

141

142

143

144

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

172

173

174

175

176

177

178

179

180

181

182

183

import vtkAll as vtk 

from vtkNumpy import addNumpyToVtk 

from shallowCopy import shallowCopy 

import numpy as np 

 

class DebugData(object): 

 

    def __init__(self): 

        self.append = vtk.vtkAppendPolyData() 

 

 

    def write(self, filename): 

        writer = vtk.vtkXMLPolyDataWriter() 

        writer.SetInputConnection(self.append.GetOutputPort()) 

        writer.SetFileName(filename) 

        writer.Update() 

 

 

    def addPolyData(self, polyData, color=[1,1,1], extraLabels=None): 

        ''' 

        Add a vtkPolyData to the debug data.  A color can be provided. 

        If the extraLabels argument is used, it should be a list of tuples, 

        each tuple is (labelName, labelValue) where labelName is a string and 

        labelValue is an int or float.  An array with labelName will be filled 

        with labelValue and added to the poly data. 

        ''' 

        polyData = shallowCopy(polyData) 

 

        if color is not None: 

            colorArray = np.empty((polyData.GetNumberOfPoints(), 3), dtype=np.uint8) 

            colorArray[:,:] = np.array(color)*255 

            addNumpyToVtk(polyData, colorArray, 'RGB255') 

 

        if extraLabels is not None: 

            for labelName, labelValue in extraLabels: 

                extraArray = np.empty((polyData.GetNumberOfPoints(), 1), dtype=type(labelValue)) 

                extraArray[:] = labelValue 

                addNumpyToVtk(polyData, extraArray, labelName) 

 

        self.append.AddInput(polyData) 

 

 

    def addLine(self, p1, p2, radius=0.0, color=[1,1,1]): 

 

        line = vtk.vtkLineSource() 

        line.SetPoint1(p1) 

        line.SetPoint2(p2) 

        line.Update() 

 

        if radius == 0.0: 

            self.addPolyData(line.GetOutput(), color) 

        else: 

            tube = vtk.vtkTubeFilter() 

            tube.SetRadius(radius) 

            tube.SetNumberOfSides(24) 

            tube.CappingOn() 

            tube.SetInputConnection(line.GetOutputPort()) 

            tube.Update() 

            self.addPolyData(tube.GetOutput(), color) 

 

    def addFrame(self, frame, scale, tubeRadius=0.0): 

 

        origin = np.array([0.0, 0.0, 0.0]) 

        axes = [[scale, 0.0, 0.0], [0.0, scale, 0.0], [0.0, 0.0, scale]] 

        colors = [[1,0,0], [0,1,0], [0,0,1]] 

        frame.TransformPoint(origin, origin) 

        for axis, color in zip(axes, colors): 

            frame.TransformVector(axis, axis) 

            self.addLine(origin, origin+axis, radius=tubeRadius, color=color) 

 

    def addCircle(self, origin, normal, radius, color=[1,1,1]): 

        self.addCone(origin, normal, radius, height=0, color=color, fill=False) 

 

    def addCone(self, origin, normal, radius, height, color=[1,1,1], fill=True): 

        cone = vtk.vtkConeSource() 

        cone.SetRadius(radius) 

        cone.SetCenter(origin) 

        cone.SetDirection(normal) 

        cone.SetHeight(height) 

        cone.SetResolution(32) 

        if fill: 

            cone.Update() 

            self.addPolyData(cone.GetOutput(), color) 

        else: 

            edges = vtk.vtkExtractEdges() 

            edges.AddInputConnection(cone.GetOutputPort()) 

            edges.Update() 

            self.addPolyData(edges.GetOutput(), color) 

 

    def addArrow(self, start, end, headRadius=0.05, tubeRadius=0.01, color=[1,1,1], startHead=False, endHead=True): 

        normal = np.array(end) - np.array(start) 

        normal = normal / np.linalg.norm(normal) 

        if startHead: 

            start = np.array(start) + headRadius * normal 

        if endHead: 

            end = np.array(end) - headRadius * normal 

        self.addLine(start, end, radius=tubeRadius, color=color) 

        if startHead: 

            self.addCone(origin=start, normal=-normal, radius=headRadius, 

                         height=headRadius, color=color, fill=True) 

        if endHead: 

            self.addCone(origin=end, normal=normal, radius=headRadius, 

                         height=headRadius, color=color, fill=True) 

 

    def addSphere(self, center, radius=0.05, color=[1,1,1], resolution=24): 

 

        sphere = vtk.vtkSphereSource() 

        sphere.SetCenter(center) 

        sphere.SetThetaResolution(resolution) 

        sphere.SetPhiResolution(resolution) 

        sphere.SetRadius(radius) 

        sphere.Update() 

        self.addPolyData(sphere.GetOutput(), color) 

 

    def addCube(self, dimensions, center, color=[1,1,1], subdivisions=0): 

 

        bmin = np.array(center) - np.array(dimensions)/2.0 

        bmax = np.array(center) + np.array(dimensions)/2.0 

        cube = vtk.vtkTessellatedBoxSource() 

        cube.SetBounds(bmin[0], bmax[0], bmin[1], bmax[1], bmin[2], bmax[2]) 

        cube.SetLevel(subdivisions) 

        cube.QuadsOn() 

        cube.Update() 

        self.addPolyData(cube.GetOutput(), color) 

 

    def addCylinder(self, center, axis, length, radius, color=[1,1,1]): 

        axis = np.asarray(axis) / np.linalg.norm(axis) 

        center = np.array(center) 

        self.addLine(center - 0.5*length*axis, center + 0.5*length*axis, radius=radius, color=color) 

 

    def addCapsule(self, center, axis, length, radius, color=[1,1,1]): 

        axis = np.asarray(axis) / np.linalg.norm(axis) 

        center = np.array(center) 

        self.addCylinder(center=center, axis=axis, radius=radius, length=length, color=color) 

        self.addSphere(center=center-0.5*length*axis, radius=radius, color=color) 

        self.addSphere(center=center+0.5*length*axis, radius=radius, color=color) 

 

    def addTorus(self, radius, thickness, resolution=30): 

 

        q = vtk.vtkSuperquadricSource() 

        q.SetToroidal(1) 

        q.SetSize(radius) 

        q.SetThetaResolution(resolution) 

        # thickness doesnt seem to match to Eucliean units. 0 is none. 1 is full. .1 is a good valve 

        q.SetThickness(thickness) 

        q.Update() 

 

        # rotate Torus so that the hole axis (internally y), is set to be z, which we use for valves 

        transform = vtk.vtkTransform() 

        transform.RotateWXYZ(90,1,0,0) 

        transformFilter=vtk.vtkTransformPolyDataFilter() 

        transformFilter.SetTransform(transform) 

        transformFilter.SetInputConnection(q.GetOutputPort()) 

        transformFilter.Update() 

        self.addPolyData(transformFilter.GetOutput()) 

 

    def addEllipsoid(self, center, radii, color=[1,1,1], alpha=1.0, resolution=24): 

        """ 

        Add an ellipsoid centered at [center] with x, y, and z principal axis radii given by 

        radii = [x_scale, y_scale, z_scale] 

        """ 

        sphere = vtk.vtkSphereSource() 

        sphere.SetCenter([0,0,0]) 

        sphere.SetThetaResolution(resolution) 

        sphere.SetPhiResolution(resolution) 

        sphere.SetRadius(1.0) 

        sphere.Update() 

 

        transform = vtk.vtkTransform() 

        transform.Translate(center) 

        transform.Scale(radii) 

 

        transformFilter = vtk.vtkTransformPolyDataFilter() 

        transformFilter.SetTransform(transform) 

        transformFilter.SetInputConnection(sphere.GetOutputPort()) 

        transformFilter.Update() 

        self.addPolyData(transformFilter.GetOutput(), color) 

 

    def getPolyData(self): 

 

        if self.append.GetNumberOfInputConnections(0): 

            self.append.Update() 

        return shallowCopy(self.append.GetOutput())