40eaf91b0426a8f7194b724ceb0b596a921f5d93
[blender-staging.git] / release / scripts / startup / bl_operators / uvcalc_lightmap.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 # <pep8 compliant>
20
21 import bpy
22 import mathutils
23
24
25 class prettyface(object):
26     __slots__ = "uv", "width", "height", "children", "xoff", "yoff", "has_parent", "rot"
27
28     def __init__(self, data):
29         self.has_parent = False
30         self.rot = False  # only used for triables
31         self.xoff = 0
32         self.yoff = 0
33
34         if type(data) == list:  # list of data
35             self.uv = None
36
37             # join the data
38             if len(data) == 2:
39                 # 2 vertical blocks
40                 data[1].xoff = data[0].width
41                 self.width = data[0].width * 2
42                 self.height = data[0].height
43
44             elif len(data) == 4:
45                 # 4 blocks all the same size
46                 d = data[0].width  # dimension x/y are the same
47
48                 data[1].xoff += d
49                 data[2].yoff += d
50
51                 data[3].xoff += d
52                 data[3].yoff += d
53
54                 self.width = self.height = d * 2
55
56             #else:
57             #    print(len(data), data)
58             #    raise "Error"
59
60             for pf in data:
61                 pf.has_parent = True
62
63             self.children = data
64
65         elif type(data) == tuple:
66             # 2 blender faces
67             # f, (len_min, len_mid, len_max)
68             self.uv = data
69
70             f1, lens1, lens1ord = data[0]
71             if data[1]:
72                 f2, lens2, lens2ord = data[1]
73                 self.width = (lens1[lens1ord[0]] + lens2[lens2ord[0]]) / 2.0
74                 self.height = (lens1[lens1ord[1]] + lens2[lens2ord[1]]) / 2.0
75             else:  # 1 tri :/
76                 self.width = lens1[0]
77                 self.height = lens1[1]
78
79             self.children = []
80
81         else:  # blender face
82             # self.uv = data.uv
83             self.uv = data.id_data.uv_textures.active.data[data.index].uv  # XXX25
84
85             # cos = [v.co for v in data]
86             cos = [data.id_data.vertices[v].co for v in data.vertices]  # XXX25
87
88             self.width = ((cos[0] - cos[1]).length + (cos[2] - cos[3]).length) / 2.0
89             self.height = ((cos[1] - cos[2]).length + (cos[0] - cos[3]).length) / 2.0
90
91             self.children = []
92
93     def spin(self):
94         if self.uv and len(self.uv) == 4:
95             self.uv = self.uv[1], self.uv[2], self.uv[3], self.uv[0]
96
97         self.width, self.height = self.height, self.width
98         self.xoff, self.yoff = self.yoff, self.xoff  # not needed?
99         self.rot = not self.rot  # only for tri pairs.
100         # print("spinning")
101         for pf in self.children:
102             pf.spin()
103
104     def place(self, xoff, yoff, xfac, yfac, margin_w, margin_h):
105         from math import pi
106
107         xoff += self.xoff
108         yoff += self.yoff
109
110         for pf in self.children:
111             pf.place(xoff, yoff, xfac, yfac, margin_w, margin_h)
112
113         uv = self.uv
114         if not uv:
115             return
116
117         x1 = xoff
118         y1 = yoff
119         x2 = xoff + self.width
120         y2 = yoff + self.height
121
122         # Scale the values
123         x1 = x1 / xfac + margin_w
124         x2 = x2 / xfac - margin_w
125         y1 = y1 / yfac + margin_h
126         y2 = y2 / yfac - margin_h
127
128         # 2 Tri pairs
129         if len(uv) == 2:
130             # match the order of angle sizes of the 3d verts with the UV angles and rotate.
131             def get_tri_angles(v1, v2, v3):
132                 a1 = (v2 - v1).angle(v3 - v1, pi)
133                 a2 = (v1 - v2).angle(v3 - v2, pi)
134                 a3 = pi - (a1 + a2)  # a3= (v2 - v3).angle(v1 - v3)
135
136                 return [(a1, 0), (a2, 1), (a3, 2)]
137
138             def set_uv(f, p1, p2, p3):
139
140                 # cos =
141                 #v1 = cos[0]-cos[1]
142                 #v2 = cos[1]-cos[2]
143                 #v3 = cos[2]-cos[0]
144
145                 # angles_co = get_tri_angles(*[v.co for v in f])
146                 angles_co = get_tri_angles(*[f.id_data.vertices[v].co for v in f.vertices])  # XXX25
147
148                 angles_co.sort()
149                 I = [i for a, i in angles_co]
150
151                 # fuv = f.uv
152                 fuv = f.id_data.uv_textures.active.data[f.index].uv  # XXX25
153
154                 if self.rot:
155                     fuv[I[2]] = p1
156                     fuv[I[1]] = p2
157                     fuv[I[0]] = p3
158                 else:
159                     fuv[I[2]] = p1
160                     fuv[I[0]] = p2
161                     fuv[I[1]] = p3
162
163             f, lens, lensord = uv[0]
164
165             set_uv(f, (x1, y1), (x1, y2 - margin_h), (x2 - margin_w, y1))
166
167             if uv[1]:
168                 f, lens, lensord = uv[1]
169                 set_uv(f, (x2, y2), (x2, y1 + margin_h), (x1 + margin_w, y2))
170
171         else:  # 1 QUAD
172             uv[1][0], uv[1][1] = x1, y1
173             uv[2][0], uv[2][1] = x1, y2
174             uv[3][0], uv[3][1] = x2, y2
175             uv[0][0], uv[0][1] = x2, y1
176
177     def __hash__(self):
178         # None unique hash
179         return self.width, self.height
180
181
182 def lightmap_uvpack(meshes,
183                       PREF_SEL_ONLY=True,
184                       PREF_NEW_UVLAYER=False,
185                       PREF_PACK_IN_ONE=False,
186                       PREF_APPLY_IMAGE=False,
187                       PREF_IMG_PX_SIZE=512,
188                       PREF_BOX_DIV=8,
189                       PREF_MARGIN_DIV=512
190                       ):
191     '''
192     BOX_DIV if the maximum division of the UV map that
193     a box may be consolidated into.
194     Basicly, a lower value will be slower but waist less space
195     and a higher value will have more clumpy boxes but more waisted space
196     '''
197     import time
198     from math import sqrt
199
200     if not meshes:
201         return
202
203     t = time.time()
204
205     if PREF_PACK_IN_ONE:
206         if PREF_APPLY_IMAGE:
207             image = bpy.data.images.new(name="lightmap", width=PREF_IMG_PX_SIZE, height=PREF_IMG_PX_SIZE, alpha=False)
208         face_groups = [[]]
209     else:
210         face_groups = []
211
212     for me in meshes:
213         # Add face UV if it does not exist.
214         # All new faces are selected.
215         if not me.uv_textures:
216             me.uv_textures.new()
217
218         if PREF_SEL_ONLY:
219             faces = [f for f in me.faces if f.select]
220         else:
221             faces = me.faces[:]
222
223         if PREF_PACK_IN_ONE:
224             face_groups[0].extend(faces)
225         else:
226             face_groups.append(faces)
227
228         if PREF_NEW_UVLAYER:
229             me.uv_textures.new()
230
231     for face_sel in face_groups:
232         print("\nStarting unwrap")
233
234         if len(face_sel) < 4:
235             print("\tWarning, less then 4 faces, skipping")
236             continue
237
238         pretty_faces = [prettyface(f) for f in face_sel if len(f.vertices) == 4]
239
240         # Do we have any tri's
241         if len(pretty_faces) != len(face_sel):
242
243             # Now add tri's, not so simple because we need to pair them up.
244             def trylens(f):
245                 # f must be a tri
246
247                 # cos = [v.co for v in f]
248                 cos = [f.id_data.vertices[v].co for v in f.vertices]  # XXX25
249
250                 lens = [(cos[0] - cos[1]).length, (cos[1] - cos[2]).length, (cos[2] - cos[0]).length]
251
252                 lens_min = lens.index(min(lens))
253                 lens_max = lens.index(max(lens))
254                 for i in range(3):
255                     if i != lens_min and i != lens_max:
256                         lens_mid = i
257                         break
258                 lens_order = lens_min, lens_mid, lens_max
259
260                 return f, lens, lens_order
261
262             tri_lengths = [trylens(f) for f in face_sel if len(f.vertices) == 3]
263             del trylens
264
265             def trilensdiff(t1, t2):
266                 return\
267                 abs(t1[1][t1[2][0]] - t2[1][t2[2][0]]) + \
268                 abs(t1[1][t1[2][1]] - t2[1][t2[2][1]]) + \
269                 abs(t1[1][t1[2][2]] - t2[1][t2[2][2]])
270
271             while tri_lengths:
272                 tri1 = tri_lengths.pop()
273
274                 if not tri_lengths:
275                     pretty_faces.append(prettyface((tri1, None)))
276                     break
277
278                 best_tri_index = -1
279                 best_tri_diff = 100000000.0
280
281                 for i, tri2 in enumerate(tri_lengths):
282                     diff = trilensdiff(tri1, tri2)
283                     if diff < best_tri_diff:
284                         best_tri_index = i
285                         best_tri_diff = diff
286
287                 pretty_faces.append(prettyface((tri1, tri_lengths.pop(best_tri_index))))
288
289         # Get the min, max and total areas
290         max_area = 0.0
291         min_area = 100000000.0
292         tot_area = 0
293         for f in face_sel:
294             area = f.area
295             if area > max_area:
296                 max_area = area
297             if area < min_area:
298                 min_area = area
299             tot_area += area
300
301         max_len = sqrt(max_area)
302         min_len = sqrt(min_area)
303         side_len = sqrt(tot_area)
304
305         # Build widths
306
307         curr_len = max_len
308
309         print("\tGenerating lengths...", end="")
310
311         lengths = []
312         while curr_len > min_len:
313             lengths.append(curr_len)
314             curr_len = curr_len / 2.0
315
316             # Dont allow boxes smaller then the margin
317             # since we contract on the margin, boxes that are smaller will create errors
318             # print(curr_len, side_len/MARGIN_DIV)
319             if curr_len / 4.0 < side_len / PREF_MARGIN_DIV:
320                 break
321
322         if not lengths:
323             lengths.append(curr_len)
324
325         # convert into ints
326         lengths_to_ints = {}
327
328         l_int = 1
329         for l in reversed(lengths):
330             lengths_to_ints[l] = l_int
331             l_int *= 2
332
333         lengths_to_ints = list(lengths_to_ints.items())
334         lengths_to_ints.sort()
335         print("done")
336
337         # apply quantized values.
338
339         for pf in pretty_faces:
340             w = pf.width
341             h = pf.height
342             bestw_diff = 1000000000.0
343             besth_diff = 1000000000.0
344             new_w = 0.0
345             new_h = 0.0
346             for l, i in lengths_to_ints:
347                 d = abs(l - w)
348                 if d < bestw_diff:
349                     bestw_diff = d
350                     new_w = i  # assign the int version
351
352                 d = abs(l - h)
353                 if d < besth_diff:
354                     besth_diff = d
355                     new_h = i  # ditto
356
357             pf.width = new_w
358             pf.height = new_h
359
360             if new_w > new_h:
361                 pf.spin()
362
363         print("...done")
364
365         # Since the boxes are sized in powers of 2, we can neatly group them into bigger squares
366         # this is done hierarchily, so that we may avoid running the pack function
367         # on many thousands of boxes, (under 1k is best) because it would get slow.
368         # Using an off and even dict us usefull because they are packed differently
369         # where w/h are the same, their packed in groups of 4
370         # where they are different they are packed in pairs
371         #
372         # After this is done an external pack func is done that packs the whole group.
373
374         print("\tConsolidating Boxes...", end="")
375         even_dict = {}  # w/h are the same, the key is an int (w)
376         odd_dict = {}  # w/h are different, the key is the (w,h)
377
378         for pf in pretty_faces:
379             w, h = pf.width, pf.height
380             if w == h:
381                 even_dict.setdefault(w, []).append(pf)
382             else:
383                 odd_dict.setdefault((w, h), []).append(pf)
384
385         # Count the number of boxes consolidated, only used for stats.
386         c = 0
387
388         # This is tricky. the total area of all packed boxes, then squt that to get an estimated size
389         # this is used then converted into out INT space so we can compare it with
390         # the ints assigned to the boxes size
391         # and divided by BOX_DIV, basicly if BOX_DIV is 8
392         # ...then the maximum box consolidataion (recursive grouping) will have a max width & height
393         # ...1/8th of the UV size.
394         # ...limiting this is needed or you end up with bug unused texture spaces
395         # ...however if its too high, boxpacking is way too slow for high poly meshes.
396         float_to_int_factor = lengths_to_ints[0][0]
397         if float_to_int_factor > 0:
398             max_int_dimension = int(((side_len / float_to_int_factor)) / PREF_BOX_DIV)
399             ok = True
400         else:
401             max_int_dimension = 0.0  # wont be used
402             ok = False
403
404         # RECURSIVE prettyface grouping
405         while ok:
406             ok = False
407
408             # Tall boxes in groups of 2
409             for d, boxes in list(odd_dict.items()):
410                 if d[1] < max_int_dimension:
411                     #\boxes.sort(key = lambda a: len(a.children))
412                     while len(boxes) >= 2:
413                         # print("foo", len(boxes))
414                         ok = True
415                         c += 1
416                         pf_parent = prettyface([boxes.pop(), boxes.pop()])
417                         pretty_faces.append(pf_parent)
418
419                         w, h = pf_parent.width, pf_parent.height
420
421                         if w > h:
422                             raise "error"
423
424                         if w == h:
425                             even_dict.setdefault(w, []).append(pf_parent)
426                         else:
427                             odd_dict.setdefault((w, h), []).append(pf_parent)
428
429             # Even boxes in groups of 4
430             for d, boxes in list(even_dict.items()):
431                 if d < max_int_dimension:
432                     boxes.sort(key=lambda a: len(a.children))
433
434                     while len(boxes) >= 4:
435                         # print("bar", len(boxes))
436                         ok = True
437                         c += 1
438
439                         pf_parent = prettyface([boxes.pop(), boxes.pop(), boxes.pop(), boxes.pop()])
440                         pretty_faces.append(pf_parent)
441                         w = pf_parent.width  # width and weight are the same
442                         even_dict.setdefault(w, []).append(pf_parent)
443
444         del even_dict
445         del odd_dict
446
447         orig = len(pretty_faces)
448
449         pretty_faces = [pf for pf in pretty_faces if not pf.has_parent]
450
451         # spin every second prettyface
452         # if there all vertical you get less efficiently used texture space
453         i = len(pretty_faces)
454         d = 0
455         while i:
456             i -= 1
457             pf = pretty_faces[i]
458             if pf.width != pf.height:
459                 d += 1
460                 if d % 2:  # only pack every second
461                     pf.spin()
462                     # pass
463
464         print("Consolidated", c, "boxes, done")
465         # print("done", orig, len(pretty_faces))
466
467         # boxes2Pack.append([islandIdx, w,h])
468         print("\tPacking Boxes", len(pretty_faces), end="...")
469         boxes2Pack = [[0.0, 0.0, pf.width, pf.height, i] for i, pf in enumerate(pretty_faces)]
470         packWidth, packHeight = mathutils.geometry.box_pack_2d(boxes2Pack)
471
472         # print(packWidth, packHeight)
473
474         packWidth = float(packWidth)
475         packHeight = float(packHeight)
476
477         margin_w = ((packWidth) / PREF_MARGIN_DIV) / packWidth
478         margin_h = ((packHeight) / PREF_MARGIN_DIV) / packHeight
479
480         # print(margin_w, margin_h)
481         print("done")
482
483         # Apply the boxes back to the UV coords.
484         print("\twriting back UVs", end="")
485         for i, box in enumerate(boxes2Pack):
486             pretty_faces[i].place(box[0], box[1], packWidth, packHeight, margin_w, margin_h)
487             # pf.place(box[1][1], box[1][2], packWidth, packHeight, margin_w, margin_h)
488         print("done")
489
490         if PREF_APPLY_IMAGE:
491             if not PREF_PACK_IN_ONE:
492                 image = Image.New("lightmap", PREF_IMG_PX_SIZE, PREF_IMG_PX_SIZE, 24)
493
494             for f in face_sel:
495                 # f.image = image
496                 f.id_data.uv_textures.active.data[f.index].image = image  # XXX25
497
498     for me in meshes:
499         me.update()
500
501     print("finished all %.2f " % (time.time() - t))
502
503     # Window.RedrawAll()
504
505
506 def unwrap(operator, context, **kwargs):
507
508     is_editmode = (bpy.context.object.mode == 'EDIT')
509     if is_editmode:
510         bpy.ops.object.mode_set(mode='OBJECT', toggle=False)
511
512     PREF_ACT_ONLY = kwargs.pop("PREF_ACT_ONLY")
513
514     meshes = []
515     if PREF_ACT_ONLY:
516         obj = context.scene.objects.active
517         if obj and obj.type == 'MESH':
518             meshes = [obj.data]
519     else:
520         meshes = {me.name: me for obj in context.selected_objects if obj.type == 'MESH' for me in (obj.data,) if not me.library if len(me.faces)}.values()
521
522     if not meshes:
523         operator.report({'ERROR'}, "No mesh object.")
524         return {'CANCELLED'}
525
526     lightmap_uvpack(meshes, **kwargs)
527
528     if is_editmode:
529         bpy.ops.object.mode_set(mode='EDIT', toggle=False)
530
531     return {'FINISHED'}
532
533 from bpy.props import BoolProperty, FloatProperty, IntProperty, EnumProperty
534
535
536 class LightMapPack(bpy.types.Operator):
537     '''Follow UVs from active quads along continuous face loops'''
538     bl_idname = "uv.lightmap_pack"
539     bl_label = "Lightmap Pack"
540     bl_options = {'REGISTER', 'UNDO'}
541
542     PREF_CONTEXT = bpy.props.EnumProperty(
543             items=(("SEL_FACES", "Selected Faces", "Space all UVs evently"),
544                    ("ALL_FACES", "All Faces", "Average space UVs edge length of each loop"),
545                    ("ALL_OBJECTS", "Selected Mesh Object", "Average space UVs edge length of each loop")
546                    ),
547             name="Selection",
548             description="")
549
550     # Image & UVs...
551     PREF_PACK_IN_ONE = BoolProperty(name="Share Tex Space", default=True, description="Objects Share texture space, map all objects into 1 uvmap")
552     PREF_NEW_UVLAYER = BoolProperty(name="New UV Layer", default=False, description="Create a new UV layer for every mesh packed")
553     PREF_APPLY_IMAGE = BoolProperty(name="New Image", default=False, description="Assign new images for every mesh (only one if shared tex space enabled)")
554     PREF_IMG_PX_SIZE = IntProperty(name="Image Size", min=64, max=5000, default=512, description="Width and Height for the new image")
555
556     # UV Packing...
557     PREF_BOX_DIV = IntProperty(name="Pack Quality", min=1, max=48, default=12, description="Pre Packing before the complex boxpack")
558     PREF_MARGIN_DIV = FloatProperty(name="Margin", min=0.001, max=1.0, default=0.1, description="Size of the margin as a division of the UV")
559
560     def execute(self, context):
561         kwargs = self.as_keywords()
562         PREF_CONTEXT = kwargs.pop("PREF_CONTEXT")
563
564         if PREF_CONTEXT == 'SEL_FACES':
565             kwargs["PREF_ACT_ONLY"] = True
566             kwargs["PREF_SEL_ONLY"] = True
567         elif PREF_CONTEXT == 'ALL_FACES':
568             kwargs["PREF_ACT_ONLY"] = True
569             kwargs["PREF_SEL_ONLY"] = False
570         elif PREF_CONTEXT == 'ALL_OBJECTS':
571             kwargs["PREF_ACT_ONLY"] = False
572             kwargs["PREF_SEL_ONLY"] = False
573         else:
574             raise Exception("invalid context")
575
576         kwargs["PREF_MARGIN_DIV"] = int(1.0 / (kwargs["PREF_MARGIN_DIV"] / 100.0))
577
578         return unwrap(self, context, **kwargs)
579
580     def invoke(self, context, event):
581         wm = context.window_manager
582         return wm.invoke_props_dialog(self)