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