Multi-Objects: UV_OT_follow_active_quads
[blender.git] / release / scripts / startup / bl_operators / uvcalc_follow_active.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 # for full docs see...
22 # https://docs.blender.org/manual/en/dev/editors/uv_image/uv/editing/unwrapping/mapping_types.html#follow-active-quads
23
24 import bpy
25 from bpy.types import Operator
26
27
28 STATUS_OK =               (1 << 0)
29 STATUS_ERR_ACTIVE_FACE =  (1 << 1)
30 STATUS_ERR_NOT_SELECTED = (1 << 2)
31 STATUS_ERR_NOT_QUAD =     (1 << 3)
32
33 def extend(obj, operator, EXTEND_MODE):
34     import bmesh
35     me = obj.data
36     # script will fail without UVs
37     if not me.uv_layers:
38         me.uv_layers.new()
39
40     bm = bmesh.from_edit_mesh(me)
41
42     f_act = bm.faces.active
43     uv_act = bm.loops.layers.uv.active
44
45     if f_act is None:
46         return STATUS_ERR_ACTIVE_FACE
47     if not f_act.select:
48         return STATUS_ERR_NOT_SELECTED
49     elif len(f_act.verts) != 4:
50         return STATUS_ERR_NOT_QUAD
51
52     faces = [f for f in bm.faces if f.select and len(f.verts) == 4]
53
54     # our own local walker
55     def walk_face_init(faces, f_act):
56         # first tag all faces True (so we don't uvmap them)
57         for f in bm.faces:
58             f.tag = True
59         # then tag faces arg False
60         for f in faces:
61             f.tag = False
62         # tag the active face True since we begin there
63         f_act.tag = True
64
65     def walk_face(f):
66         # all faces in this list must be tagged
67         f.tag = True
68         faces_a = [f]
69         faces_b = []
70
71         while faces_a:
72             for f in faces_a:
73                 for l in f.loops:
74                     l_edge = l.edge
75                     if (l_edge.is_manifold is True) and (l_edge.seam is False):
76                         l_other = l.link_loop_radial_next
77                         f_other = l_other.face
78                         if not f_other.tag:
79                             yield (f, l, f_other)
80                             f_other.tag = True
81                             faces_b.append(f_other)
82             # swap
83             faces_a, faces_b = faces_b, faces_a
84             faces_b.clear()
85
86     def walk_edgeloop(l):
87         """
88         Could make this a generic function
89         """
90         e_first = l.edge
91         e = None
92         while True:
93             e = l.edge
94             yield e
95
96             # don't step past non-manifold edges
97             if e.is_manifold:
98                 # welk around the quad and then onto the next face
99                 l = l.link_loop_radial_next
100                 if len(l.face.verts) == 4:
101                     l = l.link_loop_next.link_loop_next
102                     if l.edge is e_first:
103                         break
104                 else:
105                     break
106             else:
107                 break
108
109     def extrapolate_uv(fac,
110                        l_a_outer, l_a_inner,
111                        l_b_outer, l_b_inner):
112         l_b_inner[:] = l_a_inner
113         l_b_outer[:] = l_a_inner + ((l_a_inner - l_a_outer) * fac)
114
115     def apply_uv(f_prev, l_prev, f_next):
116         l_a = [None, None, None, None]
117         l_b = [None, None, None, None]
118
119         l_a[0] = l_prev
120         l_a[1] = l_a[0].link_loop_next
121         l_a[2] = l_a[1].link_loop_next
122         l_a[3] = l_a[2].link_loop_next
123
124         #  l_b
125         #  +-----------+
126         #  |(3)        |(2)
127         #  |           |
128         #  |l_next(0)  |(1)
129         #  +-----------+
130         #        ^
131         #  l_a   |
132         #  +-----------+
133         #  |l_prev(0)  |(1)
134         #  |    (f)    |
135         #  |(3)        |(2)
136         #  +-----------+
137         #  copy from this face to the one above.
138
139         # get the other loops
140         l_next = l_prev.link_loop_radial_next
141         if l_next.vert != l_prev.vert:
142             l_b[1] = l_next
143             l_b[0] = l_b[1].link_loop_next
144             l_b[3] = l_b[0].link_loop_next
145             l_b[2] = l_b[3].link_loop_next
146         else:
147             l_b[0] = l_next
148             l_b[1] = l_b[0].link_loop_next
149             l_b[2] = l_b[1].link_loop_next
150             l_b[3] = l_b[2].link_loop_next
151
152         l_a_uv = [l[uv_act].uv for l in l_a]
153         l_b_uv = [l[uv_act].uv for l in l_b]
154
155         if EXTEND_MODE == 'LENGTH_AVERAGE':
156             fac = edge_lengths[l_b[2].edge.index][0] / edge_lengths[l_a[1].edge.index][0]
157         elif EXTEND_MODE == 'LENGTH':
158             a0, b0, c0 = l_a[3].vert.co, l_a[0].vert.co, l_b[3].vert.co
159             a1, b1, c1 = l_a[2].vert.co, l_a[1].vert.co, l_b[2].vert.co
160
161             d1 = (a0 - b0).length + (a1 - b1).length
162             d2 = (b0 - c0).length + (b1 - c1).length
163             try:
164                 fac = d2 / d1
165             except ZeroDivisionError:
166                 fac = 1.0
167         else:
168             fac = 1.0
169
170         extrapolate_uv(fac,
171                        l_a_uv[3], l_a_uv[0],
172                        l_b_uv[3], l_b_uv[0])
173
174         extrapolate_uv(fac,
175                        l_a_uv[2], l_a_uv[1],
176                        l_b_uv[2], l_b_uv[1])
177
178     # -------------------------------------------
179     # Calculate average length per loop if needed
180
181     if EXTEND_MODE == 'LENGTH_AVERAGE':
182         bm.edges.index_update()
183         edge_lengths = [None] * len(bm.edges)
184
185         for f in faces:
186             # we know its a quad
187             l_quad = f.loops[:]
188             l_pair_a = (l_quad[0], l_quad[2])
189             l_pair_b = (l_quad[1], l_quad[3])
190
191             for l_pair in (l_pair_a, l_pair_b):
192                 if edge_lengths[l_pair[0].edge.index] is None:
193
194                     edge_length_store = [-1.0]
195                     edge_length_accum = 0.0
196                     edge_length_total = 0
197
198                     for l in l_pair:
199                         if edge_lengths[l.edge.index] is None:
200                             for e in walk_edgeloop(l):
201                                 if edge_lengths[e.index] is None:
202                                     edge_lengths[e.index] = edge_length_store
203                                     edge_length_accum += e.calc_length()
204                                     edge_length_total += 1
205
206                     edge_length_store[0] = edge_length_accum / edge_length_total
207
208     # done with average length
209     # ------------------------
210
211     walk_face_init(faces, f_act)
212     for f_triple in walk_face(f_act):
213         apply_uv(*f_triple)
214
215     bmesh.update_edit_mesh(me, False)
216     return STATUS_OK
217
218
219 def main(context, operator):
220     num_meshes = 0
221     num_errors = 0
222     status = 0
223
224     ob_list = [ob for ob in context.selected_objects if ob and ob.type == 'MESH']
225     for ob in ob_list:
226         ob.data.tag = False
227
228     for ob in ob_list:
229         if ob.data.tag:
230             continue
231
232         num_meshes += 1
233         ob.data.tag = True
234
235         ret = extend(ob, operator, operator.properties.mode)
236         if ret != STATUS_OK:
237             num_errors += 1
238             status |= ret
239
240     if num_errors == num_meshes:
241         if status & STATUS_ERR_NOT_QUAD:
242             operator.report({'ERROR'}, "Active face must be a quad")
243         elif status & STATUS_ERR_NOT_SELECTED:
244             operator.report({'ERROR'}, "Active face not selected")
245         else:
246             assert((status & STATUS_ERR_ACTIVE_FACE) != 0)
247             operator.report({'ERROR'}, "No active face")
248
249
250 class FollowActiveQuads(Operator):
251     """Follow UVs from active quads along continuous face loops"""
252     bl_idname = "uv.follow_active_quads"
253     bl_label = "Follow Active Quads"
254     bl_options = {'REGISTER', 'UNDO'}
255
256     mode: bpy.props.EnumProperty(
257         name="Edge Length Mode",
258         description="Method to space UV edge loops",
259         items=(('EVEN', "Even", "Space all UVs evenly"),
260                ('LENGTH', "Length", "Average space UVs edge length of each loop"),
261                ('LENGTH_AVERAGE', "Length Average", "Average space UVs edge length of each loop"),
262                ),
263         default='LENGTH_AVERAGE',
264     )
265
266     @classmethod
267     def poll(cls, context):
268         obj = context.active_object
269         return (obj is not None and obj.type == 'MESH')
270
271     def execute(self, context):
272         main(context, self)
273         return {'FINISHED'}
274
275     def invoke(self, context, event):
276         wm = context.window_manager
277         return wm.invoke_props_dialog(self)
278
279
280 classes = (
281     FollowActiveQuads,
282 )