e8643175353b9a476bed7b89d406db866a2f6242
[blender-addons-contrib.git] / wetted_mesh.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., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
16 #
17 # ##### END GPL LICENSE BLOCK #####
18
19 bl_info = {
20     "name": "Add Wetted Mesh",
21     "author": "freejack",
22     "version": (0, 2, 1),
23     "blender": (2, 5, 8),
24     "location": "View3D > Tool Shelf > Wetted Mesh Panel",
25     "description": "Adds separated fluid, dry and wetted mesh for selected pair.",
26     "warning": "",
27     "wiki_url": "http://wiki.blender.org/index.php/Extensions:2.5/Py/"\
28         "Scripts/Mesh/Wetted_Mesh",
29     "tracker_url": "http://projects.blender.org/tracker/index.php?"\
30         "func=detail&aid=27156",
31     "category": "Mesh"}
32
33 import bpy
34 import collections
35 import math
36
37 ### Tool Panel ###
38 class VIEW3D_PT_tools_WettedMesh(bpy.types.Panel):
39     '''Wetted Mesh Tool Panel'''
40     bl_space_type = 'VIEW_3D'
41     bl_region_type = 'TOOLS'
42     bl_label = 'Wetted Mesh'
43     bl_context = 'objectmode'
44
45     def draw(self, context):
46         layout = self.layout
47         col = layout.column(align=True)
48         slcnt = len(context.selected_objects)
49
50         if slcnt != 2:
51             col.label(text = 'Select two mesh objects')
52             col.label(text = 'to generate separated')
53             col.label(text = 'fluid, dry and wetted')
54             col.label(text = 'meshes.')
55         else:
56             (solid, fluid) = getSelectedPair(context)
57             col.label(text = 'solid = '+solid.name)
58             col.label(text = 'fluid = '+fluid.name)
59             col.operator('mesh.primitive_wetted_mesh_add', text='Generate Meshes')
60
61 ### Operator ###
62 class AddWettedMesh(bpy.types.Operator):
63     '''Add wetted mesh for selected mesh pair'''
64     bl_idname = "mesh.primitive_wetted_mesh_add"
65     bl_label = "Add Wetted Mesh"
66     bl_options = {'REGISTER', 'UNDO'}
67     statusMessage = ''
68
69     def draw(self, context):
70         layout = self.layout
71         col = layout.column(align=True)
72         col.label(text = self.statusMessage)
73
74     def execute(self, context):
75         # make sure a pair of objects is selected
76         if len(context.selected_objects) != 2:
77             # should not happen if called from tool panel
78             self.report({'WARNING'}, "no mesh pair selected, operation cancelled")
79             return {'CANCELLED'}
80
81         print("add_wetted_mesh begin")
82         
83         # super-selected object is solid, other object is fluid
84         (solid, fluid) = getSelectedPair(context)
85         print("   solid = "+solid.name)
86         print("   fluid = "+fluid.name)
87             
88         # make a copy of fluid object, convert to mesh if required
89         print("   copy fluid")
90         bpy.ops.object.select_all(action='DESELECT')
91         fluid.select = True
92         context.scene.objects.active = fluid
93         bpy.ops.object.duplicate()
94         bpy.ops.object.convert(target='MESH', keep_original=False)
95         bpy.ops.object.transform_apply(location=True, rotation=True, scale=True)
96         fluidCopy = context.object
97         
98         # substract solid from fluidCopy
99         print("   bool: fluidCopy DIFFERENCE solid")
100         bpy.ops.object.modifier_add(type='BOOLEAN')
101         bop = fluidCopy.modifiers.items()[0]
102         bop[1].operation = 'DIFFERENCE'
103         bop[1].object = solid
104         bpy.ops.object.modifier_apply(apply_as='DATA', modifier=bop[0])
105         fluidMinusSolid = fluidCopy
106         fluidMinusSolid.name = "fluidMinusSolid"
107         
108         # make a second copy of fluid object
109         print("   copy fluid")
110         bpy.ops.object.select_all(action='DESELECT')
111         fluid.select = True
112         context.scene.objects.active = fluid
113         bpy.ops.object.duplicate()
114         bpy.ops.object.convert(target='MESH', keep_original=False)
115         bpy.ops.object.transform_apply(location=True, rotation=True, scale=True)
116         fluidCopy = context.object
117         
118         # make union from fluidCopy and solid
119         print("   bool: fluidCopy UNION solid")
120         bpy.ops.object.modifier_add(type='BOOLEAN')
121         bop = fluidCopy.modifiers.items()[0]
122         bop[1].operation = 'UNION'
123         bop[1].object = solid
124         bpy.ops.object.modifier_apply(apply_as='DATA', modifier=bop[0])
125         fluidUnionSolid = fluidCopy
126         fluidUnionSolid.name = "fluidUnionSolid"
127         
128         # index meshes
129         print("   KDTree index fluidMinusSolid")
130         fluidMinusSolidKDT = KDTree(3, fluidMinusSolid.data.vertices)
131         print("   KDTree index fluidUnionSolid")
132         fluidUnionSolidKDT = KDTree(3, fluidUnionSolid.data.vertices)
133         kdtrees = (fluidMinusSolidKDT, fluidUnionSolidKDT)
134         
135         # build mesh face sets
136         faceDict = { }
137         vertDict = { }
138         
139         print("   processing fluidMinusSolid faces")
140         cacheDict = { }
141         setFMSfaces = set()
142         numFaces = len(fluidUnionSolid.data.faces)
143         i = 0
144         for f in fluidMinusSolid.data.faces:
145             if i % 500 == 0:
146                 print("      ", i, " / ", numFaces)
147             i += 1
148             fuid = unifiedFaceId(kdtrees, f, fluidMinusSolid.data.vertices, \
149                                  faceDict, vertDict, cacheDict)
150             setFMSfaces.add(fuid)
151         
152         print("   processing fluidUnionSolid faces")
153         cacheDict = { }
154         setFUSfaces = set()
155         numFaces = len(fluidUnionSolid.data.faces)
156         i = 0
157         for f in fluidUnionSolid.data.faces:
158             if i % 500 == 0:
159                 print("      ", i, " / ", numFaces)
160             i += 1
161             fuid = unifiedFaceId(kdtrees, f, fluidUnionSolid.data.vertices, \
162                                  faceDict, vertDict, cacheDict)
163             setFUSfaces.add(fuid)
164         
165         # remove boolean helpers
166         print("   delete helper objects")
167         bpy.ops.object.select_all(action='DESELECT')
168         fluidUnionSolid.select = True
169         fluidMinusSolid.select = True
170         bpy.ops.object.delete()
171
172         # wetted = FMS - FUS
173         print("   set operation FMS diff FUS")
174         setWetFaces = setFMSfaces.difference(setFUSfaces)
175         print("   build wetted mesh")
176         verts, faces = buildMesh(setWetFaces, faceDict, vertDict)
177         print("   create wetted mesh")
178         wetted = createMesh("Wetted", verts, faces)
179
180         # fluid = FMS x FUS
181         print("   set operation FMS intersect FUS")
182         setFluidFaces = setFMSfaces.intersection(setFUSfaces)
183         print("   build fluid mesh")
184         verts, faces = buildMesh(setFluidFaces, faceDict, vertDict)
185         print("   create fluid mesh")
186         fluid = createMesh("Fluid", verts, faces)
187         
188         # solid = FUS - FMS
189         print("   set operation FUS diff FMS")
190         setSolidFaces = setFUSfaces.difference(setFMSfaces)
191         print("   build solid mesh")
192         verts, faces = buildMesh(setSolidFaces, faceDict, vertDict)
193         print("   create solid mesh")
194         solid = createMesh("Solid", verts, faces)
195         
196         # parent wetted mesh
197         print("   parent mesh")
198         bpy.ops.object.add(type='EMPTY')
199         wettedMesh = context.object
200         solid.select = True
201         fluid.select = True
202         wetted.select = True
203         wettedMesh.select = True
204         bpy.ops.object.parent_set(type='OBJECT')
205         wettedMesh.name = 'WettedMesh'
206         
207         print("add_wetted_mesh done")
208         self.statusMessage = 'created '+wettedMesh.name
209
210         return {'FINISHED'}
211
212
213 ### Registration ###
214 def register():
215     bpy.utils.register_class(VIEW3D_PT_tools_WettedMesh)
216     bpy.utils.register_class(AddWettedMesh)
217
218
219 def unregister():
220     bpy.utils.unregister_class(VIEW3D_PT_tools_WettedMesh)
221     bpy.utils.unregister_class(AddWettedMesh)
222
223 if __name__ == "__main__":
224     register()
225
226
227 #
228 # KD tree (used to create a geometric index of mesh vertices)
229 #
230
231 def distance(a, b):
232     return (a-b).length
233
234 Node = collections.namedtuple("Node", 'point axis label left right')
235
236 class KDTree(object):
237     """A tree for nearest neighbor search in a k-dimensional space.
238
239     For information about the implementation, see
240     http://en.wikipedia.org/wiki/Kd-tree
241
242     Usage:
243     objects is an iterable of (co, index) tuples (so MeshVertex is useable)
244     k is the number of dimensions (=3)
245     
246     t = KDTree(k, objects)
247     point, label, distance = t.nearest_neighbor(destination)
248     """
249
250     def __init__(self, k, objects=[]):
251
252         def build_tree(objects, axis=0):
253
254             if not objects:
255                 return None
256
257             objects.sort(key=lambda o: o.co[axis])
258             median_idx = len(objects) // 2
259             median_point = objects[median_idx].co
260             median_label = objects[median_idx].index
261
262             next_axis = (axis + 1) % k
263             return Node(median_point, axis, median_label,
264                         build_tree(objects[:median_idx], next_axis),
265                         build_tree(objects[median_idx + 1:], next_axis))
266
267         self.root = build_tree(list(objects))
268         self.size = len(objects)
269
270
271     def nearest_neighbor(self, destination):
272
273         best = [None, None, float('inf')]
274         # state of search: best point found, its label,
275         # lowest distance
276
277         def recursive_search(here):
278
279             if here is None:
280                 return
281             point, axis, label, left, right = here
282
283             here_sd = distance(point, destination)
284             if here_sd < best[2]:
285                 best[:] = point, label, here_sd
286
287             diff = destination[axis] - point[axis]
288             close, away = (left, right) if diff <= 0 else (right, left)
289
290             recursive_search(close)
291             if math.fabs(diff) < best[2]:
292                 recursive_search(away)
293
294         recursive_search(self.root)
295         return best[0], best[1], best[2]
296
297
298 #
299 # helper functions
300 #
301
302 # get super-selected object and other object from selected pair
303 def getSelectedPair(context):
304     objA = context.object
305     objB = context.selected_objects[0]
306     if objA == objB:
307         objB = context.selected_objects[1]
308     return (objA, objB)
309
310 # get a unified vertex id for given coordinates
311 def unifiedVertexId(kdtrees, location, vertDict):
312     eps = 0.0001
313     offset = 0
314     for t in kdtrees:
315         co, index, d = t.nearest_neighbor(location)
316         if d < eps:
317             uvid = offset + index
318             if uvid not in vertDict:
319                 vertDict[uvid] = co
320             return uvid
321         offset += t.size
322     return -1
323
324 # get a unified face id tuple
325 #    Stores the ordered face id tuple in faceDict
326 #    and the used coordinates for vertex id in vertDict.
327 #    cacheDict caches the unified vertex id (lookup in kdtree is expensive).
328 #    For each mesh (where the face belongs to) a separate cacheDict is expected.
329 def unifiedFaceId(kdtrees, face, vertices, faceDict, vertDict, cacheDict):
330     fids = [ ]
331     for v in face.vertices:
332         uvid = cacheDict.get(v)
333         if uvid == None:
334             uvid = unifiedVertexId(kdtrees, vertices[v].co, vertDict)
335             cacheDict[v] = uvid
336         fids.append(uvid)
337     ofids = tuple(fids)
338     fids.sort()
339     fuid = tuple(fids)
340     if fuid not in faceDict:
341         faceDict[fuid] = ofids
342     return fuid
343
344 # build vertex and face array from unified face sets
345 def buildMesh(unifiedFaceSet, faceDict, vertDict):
346     verts = [ ]
347     nextV = 0
348     myV = { }
349     faces = [ ]
350     for uf in unifiedFaceSet:
351         of = faceDict[uf]
352         myf = [ ]
353         for uV in of:
354             v = myV.get(uV)
355             if v == None:
356                 v = nextV
357                 myV[uV] = nextV
358                 verts.append(vertDict[uV])
359                 nextV += 1
360             myf.append(v)
361         faces.append(myf)
362     return verts, faces
363
364 # create mesh object and link to scene
365 def createMesh(name, verts, faces):
366     me = bpy.data.meshes.new(name+"Mesh")
367     ob = bpy.data.objects.new(name, me)
368     ob.show_name = True
369     bpy.context.scene.objects.link(ob)
370     me.from_pydata(verts, [], faces)
371     me.update(calc_edges=True)
372     return ob