chiark / gitweb /
Fix dae matrix and instance_node support. Which makes sketchup exports working.
[cura.git] / Cura / util / dae.py
1 import sys, math, re, os, struct, time
2 from  xml.parsers.expat import ParserCreate
3
4 import mesh
5
6 class daeModel(mesh.mesh):
7         def __init__(self):
8                 super(daeModel, self).__init__()
9
10         def load(self, filename):
11                 r = ParserCreate()
12                 r.StartElementHandler = self._StartElementHandler
13                 r.EndElementHandler = self._EndElementHandler
14                 r.CharacterDataHandler = self._CharacterDataHandler
15
16                 self._base = {}
17                 self._cur = self._base
18                 self._idMap = {}
19                 self._geometryList = []
20                 r.ParseFile(open(filename, "r"))
21                 
22                 self.vertexCount = 0
23                 for instance_visual_scene in self._base['collada'][0]['scene'][0]['instance_visual_scene']:
24                         for node in self._idMap[instance_visual_scene['_url']]['node']:
25                                 self._ProcessNode1(node)
26                 self._prepareVertexCount(self.vertexCount)
27                 for instance_visual_scene in self._base['collada'][0]['scene'][0]['instance_visual_scene']:
28                         for node in self._idMap[instance_visual_scene['_url']]['node']:
29                                 self._ProcessNode2(node)
30                 
31                 self._base = None
32                 self._cur = None
33                 self._idMap = None
34                 
35                 self._postProcessAfterLoad()
36                 return self
37         
38         def _ProcessNode1(self, node):
39                 if 'node' in node:
40                         for n in node['node']:
41                                 self._ProcessNode1(n)
42                 if 'instance_geometry' in node:
43                         for instance_geometry in node['instance_geometry']:
44                                 mesh = self._idMap[instance_geometry['_url']]['mesh'][0]
45                                 if 'triangles' in mesh:
46                                         for triangles in mesh['triangles']:
47                                                 self.vertexCount += int(triangles['_count']) * 3
48                                 elif 'lines' in mesh:
49                                         pass #Ignore lines
50                                 else:
51                                         print mesh.keys()
52                 if 'instance_node' in node:
53                         for instance_node in node['instance_node']:
54                                 self._ProcessNode1(self._idMap[instance_node['_url']])
55
56         def _ProcessNode2(self, node, matrix = None):
57                 if 'matrix' in node:
58                         oldMatrix = matrix
59                         matrix = map(float, node['matrix'][0]['__data'].split())
60                         if oldMatrix != None:
61                                 newMatrix = [0]*16
62                                 newMatrix[0] = oldMatrix[0] * matrix[0] + oldMatrix[1] * matrix[4] + oldMatrix[2] * matrix[8] + oldMatrix[3] * matrix[12]
63                                 newMatrix[1] = oldMatrix[0] * matrix[1] + oldMatrix[1] * matrix[5] + oldMatrix[2] * matrix[9] + oldMatrix[3] * matrix[13]
64                                 newMatrix[2] = oldMatrix[0] * matrix[2] + oldMatrix[1] * matrix[6] + oldMatrix[2] * matrix[10] + oldMatrix[3] * matrix[14]
65                                 newMatrix[3] = oldMatrix[0] * matrix[3] + oldMatrix[1] * matrix[7] + oldMatrix[2] * matrix[11] + oldMatrix[3] * matrix[15]
66                                 newMatrix[4] = oldMatrix[4] * matrix[0] + oldMatrix[5] * matrix[4] + oldMatrix[6] * matrix[8] + oldMatrix[7] * matrix[12]
67                                 newMatrix[5] = oldMatrix[4] * matrix[1] + oldMatrix[5] * matrix[5] + oldMatrix[6] * matrix[9] + oldMatrix[7] * matrix[13]
68                                 newMatrix[6] = oldMatrix[4] * matrix[2] + oldMatrix[5] * matrix[6] + oldMatrix[6] * matrix[10] + oldMatrix[7] * matrix[14]
69                                 newMatrix[7] = oldMatrix[4] * matrix[3] + oldMatrix[5] * matrix[7] + oldMatrix[6] * matrix[11] + oldMatrix[7] * matrix[15]
70                                 newMatrix[8] = oldMatrix[8] * matrix[0] + oldMatrix[9] * matrix[4] + oldMatrix[10] * matrix[8] + oldMatrix[11] * matrix[12]
71                                 newMatrix[9] = oldMatrix[8] * matrix[1] + oldMatrix[9] * matrix[5] + oldMatrix[10] * matrix[9] + oldMatrix[11] * matrix[13]
72                                 newMatrix[10] = oldMatrix[8] * matrix[2] + oldMatrix[9] * matrix[6] + oldMatrix[10] * matrix[10] + oldMatrix[11] * matrix[14]
73                                 newMatrix[11] = oldMatrix[8] * matrix[3] + oldMatrix[9] * matrix[7] + oldMatrix[10] * matrix[11] + oldMatrix[11] * matrix[15]
74                                 newMatrix[12] = oldMatrix[12] * matrix[0] + oldMatrix[13] * matrix[4] + oldMatrix[14] * matrix[8] + oldMatrix[15] * matrix[12]
75                                 newMatrix[13] = oldMatrix[12] * matrix[1] + oldMatrix[13] * matrix[5] + oldMatrix[14] * matrix[9] + oldMatrix[15] * matrix[13]
76                                 newMatrix[14] = oldMatrix[12] * matrix[2] + oldMatrix[13] * matrix[6] + oldMatrix[14] * matrix[10] + oldMatrix[15] * matrix[14]
77                                 newMatrix[15] = oldMatrix[12] * matrix[3] + oldMatrix[13] * matrix[7] + oldMatrix[14] * matrix[11] + oldMatrix[15] * matrix[15]
78                                 matrix = newMatrix
79                 if 'node' in node:
80                         for n in node['node']:
81                                 self._ProcessNode2(n, matrix)
82                 if 'instance_geometry' in node:
83                         for instance_geometry in node['instance_geometry']:
84                                 mesh = self._idMap[instance_geometry['_url']]['mesh'][0]
85                                 
86                                 if 'triangles' in mesh:
87                                         for triangles in mesh['triangles']:
88                                                 for input in triangles['input']:
89                                                         if input['_semantic'] == 'VERTEX':
90                                                                 vertices = self._idMap[input['_source']]
91                                                 for input in vertices['input']:
92                                                         if input['_semantic'] == 'POSITION':
93                                                                 vertices = self._idMap[input['_source']]
94                                                 indexList = map(int, triangles['p'][0]['__data'].split())
95                                                 positionList = map(float, vertices['float_array'][0]['__data'].split())
96
97                                                 vertexCount = int(triangles['_count']) * 3
98                                                 stepSize = len(indexList) / vertexCount
99                                                 for i in xrange(0, vertexCount):
100                                                         idx = indexList[i * stepSize]
101                                                         if matrix != None:
102                                                                 x = positionList[idx*3]
103                                                                 y = positionList[idx*3+1]
104                                                                 z = positionList[idx*3+2]
105                                                                 self.addVertex(x * matrix[0] + y * matrix[1] + z * matrix[2] + matrix[3], x * matrix[4] + y * matrix[5] + z * matrix[6] + matrix[7], x * matrix[8] + y * matrix[9] + z * matrix[10] + matrix[11])
106                                                         else:
107                                                                 self.addVertex(positionList[idx*3], positionList[idx*3+1], positionList[idx*3+2])
108                 if 'instance_node' in node:
109                         for instance_node in node['instance_node']:
110                                 self._ProcessNode2(self._idMap[instance_node['_url']], matrix)
111         
112         def _StartElementHandler(self, name, attributes):
113                 name = name.lower()
114                 if not name in self._cur:
115                         self._cur[name] = []
116                 new = {'__name': name, '__parent': self._cur}
117                 self._cur[name].append(new)
118                 self._cur = new
119                 for k in attributes.keys():
120                         self._cur['_' + k] = attributes[k]
121                 
122                 if 'id' in attributes:
123                         self._idMap['#' + attributes['id']] = self._cur
124                 
125         def _EndElementHandler(self, name):
126                 self._cur = self._cur['__parent']
127
128         def _CharacterDataHandler(self, data):
129                 if len(data.strip()) < 1:
130                         return
131                 if '__data' in self._cur:
132                         self._cur['__data'] += data
133                 else:
134                         self._cur['__data'] = data
135         
136         def _GetWithKey(self, item, basename, key, value):
137                 input = basename
138                 while input in item:
139                         if item[basename]['_'+key] == value:
140                                 return self._idMap[item[input]['_source']]
141                         basename += "!"