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