Support for the C Macro system in Python.
[blender.git] / release / scripts / modules / retopo.py
1 # ##### BEGIN GPL LICENSE BLOCK #####
2 #
3 #  This program is free software; you can redistribute it and/or
4 #  modify it under the terms of the GNU General Public License
5 #  as published by the Free Software Foundation; either version 2
6 #  of the License, or (at your option) any later version.
7 #
8 #  This program is distributed in the hope that it will be useful,
9 #  but WITHOUT ANY WARRANTY; without even the implied warranty of
10 #  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
11 #  GNU General Public License for more details.
12 #
13 #  You should have received a copy of the GNU General Public License
14 #  along with this program; if not, write to the Free Software Foundation,
15 #  Inc., 59 Temple Place - Suite 330, Boston, MA  02111-1307, USA.
16 #
17 # ##### END GPL LICENSE BLOCK #####
18
19 # <pep8 compliant>
20
21 import bpy
22
23 EPS_SPLINE_DIV = 15.0 # remove doubles is ~15th the length of the spline
24
25
26 def get_hub(co, _hubs, EPS_SPLINE):
27
28     if 1:
29         for hub in _hubs.values():
30             if (hub.co - co).length < EPS_SPLINE:
31                 return hub
32
33         key = co.toTuple(3)
34         hub = _hubs[key] = Hub(co, key, len(_hubs))
35         return hub
36     else:
37         pass
38
39         '''
40         key = co.toTuple(3)
41         try:
42             return _hubs[key]
43         except:
44             hub = _hubs[key] = Hub(co, key, len(_hubs))
45             return hub
46         '''
47
48
49 class Hub(object):
50     __slots__ = "co", "key", "index", "links"
51
52     def __init__(self, co, key, index):
53         self.co = co.copy()
54         self.key = key
55         self.index = index
56         self.links = []
57
58     def get_weight(self):
59         f = 0.0
60
61         for hub_other in self.links:
62             f += (self.co - hub_other.co).length
63
64     def replace(self, other):
65         for hub in self.links:
66             try:
67                 hub.links.remove(self)
68             except:
69                 pass
70             if other not in hub.links:
71                 hub.links.append(other)
72
73     def dist(self, other):
74         return (self.co - other.co).length
75
76     def calc_faces(self, hub_ls):
77         faces = []
78         # first tris
79         for l_a in self.links:
80             for l_b in l_a.links:
81                 if l_b is not self and l_b in self.links:
82                     # will give duplicates
83                     faces.append((self.index, l_a.index, l_b.index))
84
85         # now quads, check which links share 2 different verts directly
86         def validate_quad(face):
87             if len(set(face)) != len(face):
88                 return False
89             if hub_ls[face[0]] in hub_ls[face[2]].links:
90                 return False
91             if hub_ls[face[2]] in hub_ls[face[0]].links:
92                 return False
93
94             if hub_ls[face[1]] in hub_ls[face[3]].links:
95                 return False
96             if hub_ls[face[3]] in hub_ls[face[1]].links:
97                 return False
98
99             return True
100
101         for i, l_a in enumerate(self.links):
102             links_a = set([l.index for l in l_a.links])
103             for j in range(i):
104                 l_b = self.links[j]
105
106                 links_b = set([l.index for l in l_b.links])
107
108                 isect = links_a.intersection(links_b)
109                 if len(isect) == 2:
110                     isect = list(isect)
111
112                     # check there are no diagonal lines
113                     face = (isect[0], l_a.index, isect[1], l_b.index)
114                     if validate_quad(face):
115
116                         faces.append(face)
117
118         return faces
119
120
121 class Spline:
122     __slots__ = "points", "hubs", "length"
123
124     def __init__(self, points):
125         self.points = points
126         self.hubs = []
127
128         # calc length
129         f = 0.0
130         co_prev = self.points[0]
131         for co in self.points[1:]:
132             f += (co - co_prev).length
133             co_prev = co
134         self.length = f
135
136     def link(self):
137         if len(self.hubs) < 2:
138             return
139
140         edges = list(set([i for i, hub in self.hubs]))
141         edges.sort()
142
143         edges_order = {}
144         for i in edges:
145             edges_order[i] = []
146
147
148         # self.hubs.sort()
149         for i, hub in self.hubs:
150             edges_order[i].append(hub)
151
152         hubs_order = []
153         for i in edges:
154             ls = edges_order[i]
155             edge_start = self.points[i]
156             ls.sort(key=lambda hub: (hub.co - edge_start).length)
157             hubs_order.extend(ls)
158
159         # Now we have the order, connect the hubs
160         hub_prev = hubs_order[0]
161
162         for hub in hubs_order[1:]:
163             hub.links.append(hub_prev)
164             hub_prev.links.append(hub)
165             hub_prev = hub
166
167
168 def get_points(stroke):
169     return [point.co.copy() for point in stroke.points]
170
171
172 def get_splines(gp):
173     for l in gp.layers:
174         if l.active: # XXX - should be layers.active
175             break
176
177     frame = l.active_frame
178
179     return [Spline(get_points(stroke)) for stroke in frame.strokes]
180
181
182 def xsect_spline(sp_a, sp_b, _hubs):
183     from Mathutils import LineIntersect
184     from Mathutils import MidpointVecs
185     from Geometry import ClosestPointOnLine
186     pt_a_prev = pt_b_prev = None
187     EPS_SPLINE = (sp_a.length + sp_b.length) / (EPS_SPLINE_DIV * 2)
188     pt_a_prev = sp_a.points[0]
189     for a, pt_a in enumerate(sp_a.points[1:]):
190         pt_b_prev = sp_b.points[0]
191         for b, pt_b in enumerate(sp_b.points[1:]):
192
193             # Now we have 2 edges
194             # print(pt_a, pt_a_prev, pt_b, pt_b_prev)
195             xsect = LineIntersect(pt_a, pt_a_prev, pt_b, pt_b_prev)
196             if xsect is not None:
197                 if (xsect[0] - xsect[1]).length <= EPS_SPLINE:
198                     f = ClosestPointOnLine(xsect[1], pt_a, pt_a_prev)[1]
199                     # if f >= 0.0-EPS_SPLINE and f <= 1.0+EPS_SPLINE: # for some reason doesnt work so well, same below
200                     if f >= 0.0 and f <= 1.0:
201                         f = ClosestPointOnLine(xsect[0], pt_b, pt_b_prev)[1]
202                         # if f >= 0.0-EPS_SPLINE and f <= 1.0+EPS_SPLINE:
203                         if f >= 0.0 and f <= 1.0:
204                             # This wont happen often
205                             co = MidpointVecs(xsect[0], xsect[1])
206                             hub = get_hub(co, _hubs, EPS_SPLINE)
207
208                             sp_a.hubs.append((a, hub))
209                             sp_b.hubs.append((b, hub))
210
211             pt_b_prev = pt_b
212
213         pt_a_prev = pt_a
214
215
216 def calculate(gp):
217     splines = get_splines(gp)
218     _hubs = {}
219
220     for i, sp in enumerate(splines):
221         for j, sp_other in enumerate(splines):
222             if j <= i:
223                 continue
224
225             xsect_spline(sp, sp_other, _hubs)
226
227     for sp in splines:
228         sp.link()
229
230     # remove these
231     hubs_ls = [hub for hub in _hubs.values() if hub.index != -1]
232
233     _hubs.clear()
234     _hubs = None
235
236     for i, hub in enumerate(hubs_ls):
237         hub.index = i
238
239     # Now we have connected hubs, write all edges!
240     def order(i1, i2):
241         if i1 > i2:
242             return i2, i1
243         return i1, i2
244
245     edges = {}
246
247     for hub in hubs_ls:
248         i1 = hub.index
249         for hub_other in hub.links:
250             i2 = hub_other.index
251             edges[order(i1, i2)] = None
252
253     verts = []
254     edges = edges.keys()
255     faces = []
256
257     for hub in hubs_ls:
258         verts.append(hub.co)
259         faces.extend(hub.calc_faces(hubs_ls))
260
261     # remove double faces
262     faces = dict([(tuple(sorted(f)), f) for f in faces]).values()
263
264     mesh = bpy.data.add_mesh("Retopo")
265     mesh.from_pydata(verts, [], faces)
266
267     scene = bpy.context.scene
268     mesh.update()
269     obj_new = bpy.data.add_object('MESH', "Torus")
270     obj_new.data = mesh
271     scene.objects.link(obj_new)
272
273     return obj_new
274
275
276 def main():
277     scene = bpy.context.scene
278     obj = bpy.context.object
279
280     gp = None
281
282     if obj:
283         gp = obj.grease_pencil
284
285     if not gp:
286         gp = scene.grease_pencil
287
288     if not gp:
289         raise Exception("no active grease pencil")
290
291     obj_new = calculate(gp)
292
293     scene.objects.active = obj_new
294     obj_new.selected = True
295
296     # nasty, recalc normals
297     bpy.ops.object.mode_set(mode='EDIT', toggle=False)
298     bpy.ops.mesh.normals_make_consistent(inside=False)
299     bpy.ops.object.mode_set(mode='OBJECT', toggle=False)