fix [#33332] UV follow active quads

rewrite the script to use bmesh connectivity info.

1 files changed, 104 insertions, 175 deletions

diff --git a/release/scripts/startup/bl_operators/uvcalc_follow_active.py b/release/scripts/startup/bl_operators/uvcalc_follow_active.py
index b60b5257984..727c4ad739f 100644
--- a/release/scripts/startup/bl_operators/uvcalc_follow_active.py
+++ b/release/scripts/startup/bl_operators/uvcalc_follow_active.py
@@ -26,195 +26,124 @@ from bpy.types import Operator
 
 
 def extend(obj, operator, EXTEND_MODE):
-    from bpy_extras import mesh_utils
-
+    import bmesh
     me = obj.data
-    me_verts = me.vertices
-
     # script will fail without UVs
     if not me.uv_textures:
         me.uv_textures.new()
+    
+    bm = bmesh.from_edit_mesh(me)
+    
+    f_act = bm.faces.active
+    uv_act = bm.loops.layers.uv.active
+    
+    if f_act is None:
+        operator.report({'ERROR'}, "No active face")
+        return
+    elif len(f_act.verts) != 4:
+        operator.report({'ERROR'}, "Active face must be a quad")
+        return
 
-    # Toggle Edit mode
-    is_editmode = (obj.mode == 'EDIT')
-    if is_editmode:
-        bpy.ops.object.mode_set(mode='OBJECT')
-
-    #t = sys.time()
-    edge_average_lengths = {}
-
-    OTHER_INDEX = 2, 3, 0, 1
-
-    def extend_uvs(face_source, face_target, edge_key):
-        """
-        Takes 2 faces,
-        Projects its extends its UV coords onto the face next to it.
-        Both faces must share an edge
-        """
-
-        def face_edge_vs(vi):
-            vlen = len(vi)
-            return [(vi[i], vi[(i + 1) % vlen]) for i in range(vlen)]
-
-        vidx_source = face_source.vertices
-        vidx_target = face_target.vertices
-
-        uv_layer = me.uv_layers.active.data
-        uvs_source = [uv_layer[i].uv for i in face_source.loop_indices]
-        uvs_target = [uv_layer[i].uv for i in face_target.loop_indices]
-
-        # vertex index is the key, uv is the value
-
-        uvs_vhash_source = {vindex: uvs_source[i] for i, vindex in enumerate(vidx_source)}
-
-        uvs_vhash_target = {vindex: uvs_target[i] for i, vindex in enumerate(vidx_target)}
-
-        edge_idxs_source = face_edge_vs(vidx_source)
-        edge_idxs_target = face_edge_vs(vidx_target)
-
-        source_matching_edge = -1
-        target_matching_edge = -1
-
-        edge_key_swap = edge_key[1], edge_key[0]
-
-        try:
-            source_matching_edge = edge_idxs_source.index(edge_key)
-        except:
-            source_matching_edge = edge_idxs_source.index(edge_key_swap)
-        try:
-            target_matching_edge = edge_idxs_target.index(edge_key)
-        except:
-            target_matching_edge = edge_idxs_target.index(edge_key_swap)
-
-        edgepair_inner_source = edge_idxs_source[source_matching_edge]
-        edgepair_inner_target = edge_idxs_target[target_matching_edge]
-        edgepair_outer_source = edge_idxs_source[OTHER_INDEX[source_matching_edge]]
-        edgepair_outer_target = edge_idxs_target[OTHER_INDEX[target_matching_edge]]
-
-        if edge_idxs_source[source_matching_edge] == edge_idxs_target[target_matching_edge]:
-            iA = 0  # Flipped, most common
-            iB = 1
-        else:  # The normals of these faces must be different
-            iA = 1
-            iB = 0
+    faces = [f for f in bm.faces if f.select and len(f.verts) == 4]
+
+    for f in faces:
+        f.tag = False
+    f_act.tag = True
+
+
+    # our own local walker
+    def walk_face(f):
+        # all faces in this list must be tagged
+        f.tag = True
+        faces_a = [f]
+        faces_b = []
+
+        while faces_a:
+            for f in faces_a:
+                for l in f.loops:
+                    l_edge = l.edge
+                    if (l_edge.is_manifold is True) and (l_edge.seam is False):
+                        l_other = l.link_loop_radial_next
+                        f_other = l_other.face
+                        if not f_other.tag:
+                            yield (f, l, f_other)
+                            f_other.tag = True
+                            faces_b.append(f_other)
+            # swap
+            faces_a, faces_b = faces_b, faces_a
+            faces_b.clear()
+
+    def extrapolate_uv(fac,
+                       l_a_outer, l_a_inner,
+                       l_b_outer, l_b_inner):
+        l_b_inner[:] = l_a_inner
+        l_b_outer[:] = l_a_inner + ((l_a_inner - l_a_outer) * fac)
+
+    def apply_uv(f_prev, l_prev, f_next):
+        l_a = [None, None, None, None]
+        l_b = [None, None, None, None]
+        
+        l_a[0] = l_prev
+        l_a[1] = l_a[0].link_loop_next
+        l_a[2] = l_a[1].link_loop_next
+        l_a[3] = l_a[2].link_loop_next
+
+        #  l_b
+        #  +-----------+
+        #  |(3)        |(2)
+        #  |           |
+        #  |l_next(0)  |(1)
+        #  +-----------+
+        #        ^
+        #  l_a   |
+        #  +-----------+
+        #  |l_prev(0)  |(1)
+        #  |    (f)    |
+        #  |(3)        |(2)
+        #  +-----------+
+        #  copy from this face to the one above.
+
+        # get the other loops 
+        l_next = l_prev.link_loop_radial_next
+        if l_next.vert != l_prev.vert:
+            l_b[1] = l_next
+            l_b[0] = l_b[1].link_loop_next
+            l_b[3] = l_b[0].link_loop_next
+            l_b[2] = l_b[3].link_loop_next
+        else:
+            l_b[0] = l_next
+            l_b[1] = l_b[0].link_loop_next
+            l_b[2] = l_b[1].link_loop_next
+            l_b[3] = l_b[2].link_loop_next
 
-        # Set the target UV's touching source face, no tricky calculations needed,
-        uvs_vhash_target[edgepair_inner_target[0]][:] = uvs_vhash_source[edgepair_inner_source[iA]]
-        uvs_vhash_target[edgepair_inner_target[1]][:] = uvs_vhash_source[edgepair_inner_source[iB]]
+        l_a_uv = [l[uv_act].uv for l in l_a]
+        l_b_uv = [l[uv_act].uv for l in l_b]
 
-        # Set the 2 UV's on the target face that are not touching
-        # for this we need to do basic expanding on the source faces UV's
         if EXTEND_MODE == 'LENGTH':
+            a0, b0, c0 = l_a[3].vert.co, l_a[0].vert.co, l_b[3].vert.co
+            a1, b1, c1 = l_a[2].vert.co, l_a[1].vert.co, l_b[2].vert.co
 
-            try:  # divide by zero is possible
-                '''
-                measure the length of each face from the middle of each edge to the opposite
-                along the axis we are copying, use this
-                '''
-                i1a = edgepair_outer_target[iB]
-                i2a = edgepair_inner_target[iA]
-                if i1a > i2a:
-                    i1a, i2a = i2a, i1a
-
-                i1b = edgepair_outer_source[iB]
-                i2b = edgepair_inner_source[iA]
-                if i1b > i2b:
-                    i1b, i2b = i2b, i1b
-                # print edge_average_lengths
-                factor = edge_average_lengths[i1a, i2a][0] / edge_average_lengths[i1b, i2b][0]
-            except:
-                # Div By Zero?
-                factor = 1.0
-
-            uvs_vhash_target[edgepair_outer_target[iB]][:] = uvs_vhash_source[edgepair_inner_source[0]] + factor * (uvs_vhash_source[edgepair_inner_source[0]] - uvs_vhash_source[edgepair_outer_source[1]])
-            uvs_vhash_target[edgepair_outer_target[iA]][:] = uvs_vhash_source[edgepair_inner_source[1]] + factor * (uvs_vhash_source[edgepair_inner_source[1]] - uvs_vhash_source[edgepair_outer_source[0]])
-
+            d1 = (a0 - b0).length + (a1 - b1).length
+            d2 = (b0 - c0).length + (b1 - c1).length
+            try:
+                fac = d2 / d1
+            except ZeroDivisionError:
+                fac = 1.0
         else:
-            # same as above but with no factors
-            uvs_vhash_target[edgepair_outer_target[iB]][:] = uvs_vhash_source[edgepair_inner_source[0]] + (uvs_vhash_source[edgepair_inner_source[0]] - uvs_vhash_source[edgepair_outer_source[1]])
-            uvs_vhash_target[edgepair_outer_target[iA]][:] = uvs_vhash_source[edgepair_inner_source[1]] + (uvs_vhash_source[edgepair_inner_source[1]] - uvs_vhash_source[edgepair_outer_source[0]])
+            fac = 1.0
 
-    face_act = me.polygons.active
-    if face_act == -1:
-        operator.report({'ERROR'}, "No active face")
-        return
+        extrapolate_uv(fac,
+                       l_a_uv[3], l_a_uv[0],
+                       l_b_uv[3], l_b_uv[0])
 
-    face_sel = [f for f in me.polygons if len(f.vertices) == 4 and f.select]
+        extrapolate_uv(fac,
+                       l_a_uv[2], l_a_uv[1],
+                       l_b_uv[2], l_b_uv[1])
 
-    face_act_local_index = -1
-    for i, f in enumerate(face_sel):
-        if f.index == face_act:
-            face_act_local_index = i
-            break
+    for f_triple in walk_face(f_act):
+        apply_uv(*f_triple)
 
-    if face_act_local_index == -1:
-        operator.report({'ERROR'}, "Active face not selected")
-        return
-
-    # Modes
-    # 0 not yet searched for.
-    # 1:mapped, use search from this face - removed!
-    # 2:all siblings have been searched. don't search again.
-    face_modes = [0] * len(face_sel)
-    face_modes[face_act_local_index] = 1  # extend UV's from this face.
-
-    # Edge connectivity
-    edge_faces = {}
-    for i, f in enumerate(face_sel):
-        for edkey in f.edge_keys:
-            try:
-                edge_faces[edkey].append(i)
-            except:
-                edge_faces[edkey] = [i]
-
-    if EXTEND_MODE == 'LENGTH':
-        edge_loops = mesh_utils.edge_loops_from_tessfaces(me, face_sel, [ed.key for ed in me.edges if ed.use_seam])
-        me_verts = me.vertices
-        for loop in edge_loops:
-            looplen = [0.0]
-            for ed in loop:
-                edge_average_lengths[ed] = looplen
-                looplen[0] += (me_verts[ed[0]].co - me_verts[ed[1]].co).length
-            looplen[0] = looplen[0] / len(loop)
-
-    # remove seams, so we don't map across seams.
-    for ed in me.edges:
-        if ed.use_seam:
-            # remove the edge pair if we can
-            try:
-                del edge_faces[ed.key]
-            except:
-                pass
-    # Done finding seams
-
-    # face connectivity - faces around each face
-    # only store a list of indices for each face.
-    face_faces = [[] for i in range(len(face_sel))]
-
-    for edge_key, faces in edge_faces.items():
-        if len(faces) == 2:  # Only do edges with 2 face users for now
-            face_faces[faces[0]].append((faces[1], edge_key))
-            face_faces[faces[1]].append((faces[0], edge_key))
-
-    # Now we know what face is connected to what other face, map them by connectivity
-    ok = True
-    while ok:
-        ok = False
-        for i in range(len(face_sel)):
-            if face_modes[i] == 1:  # searchable
-                for f_sibling, edge_key in face_faces[i]:
-                    if face_modes[f_sibling] == 0:
-                        face_modes[f_sibling] = 1  # mapped and search from.
-                        extend_uvs(face_sel[i], face_sel[f_sibling], edge_key)
-                        face_modes[i] = 1  # we can map from this one now.
-                        ok = True  # keep searching
-
-                face_modes[i] = 2  # don't search again
-
-    if is_editmode:
-        bpy.ops.object.mode_set(mode='EDIT')
-    else:
-        me.update_tag()
+    bmesh.update_edit_mesh(me, False)
 
 
 def main(context, operator):