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