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