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