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