Blender-TexTools/op_island_straighten_edge_loops.py

import bpy
import bmesh
import operator
import math
from mathutils import Vector
from collections import defaultdict
from math import pi

from . import utilities_uv


class op(bpy.types.Operator):
    bl_idname = "uv.textools_island_straighten_edge_loops"
    bl_label = "Straight edge loops"
    bl_description = "Straighten edge loops of UV Island and relax rest"
    bl_options = {'REGISTER', 'UNDO'}

    @classmethod
    def poll(cls, context):

        if not bpy.context.active_object:
            return False

        if bpy.context.active_object.type != 'MESH':
            return False

        # Only in Edit mode
        if bpy.context.active_object.mode != 'EDIT':
            return False

        # Only in UV editor mode
        if bpy.context.area.type != 'IMAGE_EDITOR':
            return False

        # Requires UV map
        if not bpy.context.object.data.uv_layers:
            return False

        if bpy.context.scene.tool_settings.uv_select_mode != 'EDGE':
            return False

        return True

    def execute(self, context):

        main(context)
        return {'FINISHED'}


def main(context):
    print("____________________________")

    # Store selection
    utilities_uv.selection_store()

    bm = bmesh.from_edit_mesh(bpy.context.active_object.data)
    uv_layers = bm.loops.layers.uv.verify()

    edges = utilities_uv.get_selected_uv_edges(bm, uv_layers)
    islands = utilities_uv.getSelectionIslands()
    uvs = utilities_uv.get_selected_uvs(bm, uv_layers)
    faces = [f for island in islands for f in island]

    # Get island faces

    # utilities_uv.selection_restore(bm, uv_layers)

    groups = get_edge_groups(bm, uv_layers, faces, edges, uvs)

    bpy.ops.mesh.select_mode(use_extend=False, use_expand=False, type='FACE')
    bpy.ops.mesh.select_all(action='DESELECT')
    for face in faces:
        face.select = True

    print("Edges {}x".format(len(edges)))
    print("Groups {}x".format(len(groups)))

    # Restore 3D face selection

    # Restore UV seams and clear pins
    bpy.ops.uv.seams_from_islands()
    bpy.ops.uv.pin(clear=True)

    edge_sets = []
    for edges in groups:
        edge_sets.append(EdgeSet(bm, uv_layers, edges, faces))
        # straighten_edges(bm, uv_layers, edges, faces)

    sorted_sets = sorted(edge_sets, key=lambda x: x.length, reverse=True)

    for edge_set in sorted_sets:
        edge_set.straighten()

    # Restore selection
    utilities_uv.selection_restore()


class EdgeSet:
    bm = None
    edges = []
    faces = []
    uv_layers = ''
    vert_to_uv = {}
    edge_length = {}
    length = 0

    def __init__(self, bm, uv_layers, edges, faces):
        self.bm = bm
        self.uv_layers = uv_layers
        self.edges = edges
        self.faces = faces

        # Get Vert to UV within faces
        self.vert_to_uv = utilities_uv.get_vert_to_uv(bm, uv_layers)

        # Get edge lengths
        self.edge_length = {}
        self.length = 0
        for e in edges:
            uv1 = self.vert_to_uv[e.verts[0]][0].uv
            uv2 = self.vert_to_uv[e.verts[1]][0].uv
            self.edge_length[e] = (uv2 - uv1).length
            self.length += self.edge_length[e]

    def straighten(self):
        print("Straight {}x at {:.2f} length ".format(
            len(self.edges), self.length))

        # Get edge angles in UV space
        angles = {}
        for edge in self.edges:
            uv1 = self.vert_to_uv[edge.verts[0]][0].uv
            uv2 = self.vert_to_uv[edge.verts[1]][0].uv
            delta = uv2 - uv1
            angle = math.atan2(delta.y, delta.x) % (math.pi/2)
            if angle >= (math.pi/4):
                angle = angle - (math.pi/2)
            angles[edge] = abs(angle)
            # print("Angle {:.2f} degr".format(angle * 180 / math.pi))

        # Pick edge with least rotation offset to U or V axis
        edge_main = sorted(angles.items(), key=operator.itemgetter(1))[0][0]

        print("Main edge: {} at {:.2f} degr".format(
            edge_main.index, angles[edge_main] * 180 / math.pi))

        # Rotate main edge to closest axis
        uvs = [uv for v in edge_main.verts for uv in self.vert_to_uv[v]]
        bpy.ops.uv.select_all(action='DESELECT')
        for uv in uvs:
            uv.select = True
        uv1 = self.vert_to_uv[edge_main.verts[0]][0].uv
        uv2 = self.vert_to_uv[edge_main.verts[1]][0].uv
        diff = uv2 - uv1
        angle = math.atan2(diff.y, diff.x) % (math.pi/2)
        if angle >= (math.pi/4):
            angle = angle - (math.pi/2)
        bpy.ops.uv.cursor_set(location=uv1 + diff/2)
        bpy.ops.transform.rotate(value=angle, orient_axis='Z', constraint_axis=(
            False, False, False), orient_type='GLOBAL', mirror=False, use_proportional_edit=False)

        # Expand edges and straighten
        count = len(self.edges)
        processed = [edge_main]
        for i in range(count):
            if(len(processed) < len(self.edges)):
                verts = set([v for e in processed for v in e.verts])
                edges_expand = [e for e in self.edges if e not in processed and (
                    e.verts[0] in verts or e.verts[1] in verts)]
                verts_ends = [
                    v for e in edges_expand for v in e.verts if v in verts]

                print("Step, proc {} exp: {}".format(
                    [e.index for e in processed], [e.index for e in edges_expand]))

                if len(edges_expand) == 0:
                    continue

                for edge in edges_expand:
                    # if edge.verts[0] in verts_ends and edge.verts[1] in verts_ends:
                    #     print("Cancel at edge {}".format(edge.index))
                    #     return

                    print("  E {} verts {} verts end: {}".format(edge.index, [
                          v.index for v in edge.verts], [v.index for v in verts_ends]))
                    v1 = [v for v in edge.verts if v in verts_ends][0]
                    v2 = [v for v in edge.verts if v not in verts_ends][0]
                    # direction
                    previous_edge = [e for e in processed if e.verts[0]
                                     in edge.verts or e.verts[1] in edge.verts][0]
                    prev_v1 = [v for v in previous_edge.verts if v != v1][0]
                    prev_v2 = [v for v in previous_edge.verts if v == v1][0]
                    direction = (
                        self.vert_to_uv[prev_v2][0].uv - self.vert_to_uv[prev_v1][0].uv).normalized()

                    for uv in self.vert_to_uv[v2]:
                        uv.uv = self.vert_to_uv[v1][0].uv + \
                            direction * self.edge_length[edge]

                print("Procesed {}x Expand {}x".format(
                    len(processed), len(edges_expand)))
                print("verts_ends: {}x".format(len(verts_ends)))
                processed.extend(edges_expand)

        # Select edges
        uvs = list(
            set([uv for e in self.edges for v in e.verts for uv in self.vert_to_uv[v]]))
        bpy.ops.uv.select_all(action='DESELECT')
        for uv in uvs:
            uv.select = True

        # Pin UV's
        bpy.ops.uv.pin()
        bpy.ops.uv.unwrap(method='ANGLE_BASED', margin=0.001)
        bpy.ops.uv.pin(clear=True)


def get_edge_groups(bm, uv_layers, faces, edges, uvs):
    print("Get edge groups, edges {}x".format(len(edges))+"x")

    bpy.ops.mesh.select_mode(use_extend=False, use_expand=False, type='EDGE')

    unmatched = edges.copy()

    groups = []

    for edge in edges:
        if edge in unmatched:

            # Loop select edge
            bpy.ops.mesh.select_all(action='DESELECT')
            edge.select = True
            bpy.ops.mesh.loop_multi_select(ring=False)

            # Isolate group within edges
            group = [e for e in bm.edges if e.select and e in edges]
            groups.append(group)

            # Remove from unmatched
            for e in group:
                if e in unmatched:
                    unmatched.remove(e)

            print("  Edge {} : Group: {}x , unmatched: {}".format(
                edge.index, len(group), len(unmatched)))

            # return
            # group = [edge]
            # for e in bm.edges:
            #     if e.select and e in unmatched:
            #         unmatched.remove(e)
            #         group.append(edge)

    return groups


bpy.utils.register_class(op)