MelUVStacker.py

# ##### BEGIN GPL LICENSE BLOCK #####
#
#  This program is free software; you can redistribute it and/or
#  modify it under the terms of the GNU General Public License
#  as published by the Free Software Foundation; either version 2
#  of the License, or (at your option) any later version.
#
#  This program is distributed in the hope that it will be useful,
#  but WITHOUT ANY WARRANTY; without even the implied warranty of
#  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#  GNU General Public License for more details.
#
#  You should have received a copy of the GNU General Public License
#  along with this program; if not, write to the Free Software Foundation,
#  Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
#
# ##### END GPL LICENSE BLOCK #####

bl_info = {
    "name": "Melsov UV Stacker",
    "author": "Melsov",
    "version": (1, 0),
    "blender": (2, 80, 0),
    "location": "View3D > Sidebar > MelUVStacker",
    "description": "Automation for stacking UV faces on top of each other",
    "warning": "",
    "wiki_url": "",
    "category": "3D View"}

import bpy
import bmesh

class MelStackParams(bpy.types.PropertyGroup):
    """todo description"""
    tolerance : bpy.props.FloatProperty(
        name="tolerance name",
        description="factor that determines how different two edge lengths can be and still be considered equal",
        default=.05,
        max=1.0,
        min=0,
        subtype="FACTOR",
        )
    epsilon : bpy.props.FloatProperty(
        name="epsilon",
        description="this may help if you need small edge lengths to positively match. if zero or less, does nothing",
        default=.0001,
        max=1.0,
        min=-1.0,
        subtype="FACTOR",
        )
   

class FaceMatchingInfo:
    def __init__(self, face ) -> None:
        self.face = face
        self.compare_start_index = 0
        self.edge_lengths = []

    def __str__(self) -> str:
        return F"FMI: {self.edge_lengths} compare_start_index : {self.compare_start_index}"

    def length(self) -> int:
        return len(self.edge_lengths)

    def _is_match(self, edge, other_edge) -> float:
        eps = bpy.context.scene.stack_params.epsilon
        if edge < eps and other_edge < eps:
            return 1.0
        dif = edge - other_edge         
        tolerance = edge * bpy.context.scene.stack_params.tolerance 
        if abs(dif) > tolerance:
            return -1.0
        return 1.0 - abs(dif) / tolerance

    def rate_match_with(self, other) -> float:
        if self.length() != other.length(): 
            return -1.0
        for i in range(0, self.length()):
            r = self._rate_match_with(other, i)
            if r > 0.0:
                other.compare_start_index = i
                return r
        return -1.0

    def _rate_match_with(self, other, other_offset : int) -> float:
        total = 0.0
        for i in range(0, self.length()):
            other_idx = i # (other.length() - 1 -i) if _is_mirror else i
            m = self._is_match(self.edge_lengths[i], other.edge_lengths[(other_idx + other_offset) % other.length()])
            if m < 0.0:
                return -1.0
            total += m
        return total / self.length()

class MatchedFaces:
    def __init__(self) -> None:
        self.fmis : list[FaceMatchingInfo] = []

    def getReferenceFMI(self) -> FaceMatchingInfo:
        if(len(self.fmis) == 0):
            return None
        return self.fmis[0]
    
    def addCandidate(self, fmi : FaceMatchingInfo) -> bool :
        refFMI = self.getReferenceFMI()
        if self.getReferenceFMI() is None:
            self.fmis.append(fmi)
            return True
        if refFMI is fmi:
            return False
        match_rating = refFMI.rate_match_with(fmi)
        if match_rating > 0.0:
            self.fmis.append(fmi)
            return True
        return False
    
    def length(self) -> int:
        return len(self.fmis)

class FacesLookup:
    def __init__(self) -> None:
        self.lookup = {}

    def __repr__(self):
        return F"FLookup()"
    def __str__(self):
        return F"FaceLookup: {[v.__str__()  for v in self.lookup.values()]}"

class MelUVStacker(bpy.types.Operator):
    """Mel's UV Stacker"""
    bl_idname = "object.uv_stacker"
    bl_label = "Mel UV Stacker"

    custom_key : bpy.props.StringProperty()
    
    def bmeshFrom(self, mesh):
        if mesh.is_editmode:
            return bmesh.from_edit_mesh(mesh)
        bm = bmesh.new()
        bm.from_mesh(mesh)
        return bm

    def apply_bmesh_to_mesh(self, bm, mesh) -> None:
        if bpy.context.active_object.mode == 'OBJECT':
            bm.to_mesh(mesh)
        else:
            bmesh.update_edit_mesh(mesh)

    def __init__(self) -> None:
        super().__init__()
        self._faces : FacesLookup = FacesLookup()
        self._matched_faceses : list[MatchedFaces] = []

    def _matchFace(self, fmi : FaceMatchingInfo) -> bool:
        for matches in self._matched_faceses:
            if matches.addCandidate(fmi):
                return True
        return False
  
    def _sets_from_selected(self, bm):
        for fmi in self._faces.lookup.values():
            if fmi.face.select:
                matched = MatchedFaces()
                matched.addCandidate(fmi)
                self._matched_faceses.append(matched)

    def _addMatchedFaceSets(self):
        # just throw together with the first fmi that matches
        for candidate in self._faces.lookup.values():
            if self._matchFace(candidate):
                continue
            # we failed to find a match with an existing face set
            # make a new one
            matched = MatchedFaces()
            matched.addCandidate(candidate)
            self._matched_faceses.append(matched)

    def _onlyMatchWithFaceSets(self):
        # try to add candidate faces to face sets
        #  don't make new sets if matching fails
        for candidate in self._faces.lookup.values():
            self._matchFace(candidate)
                
    def _smart_project(self):
        #  if we're in object mode switch to edit mode for this op
        #  then switch back if we switched 
        was_object_mode = bpy.context.active_object.mode == 'OBJECT'
        bpy.ops.object.mode_set(mode='EDIT')
        bpy.ops.uv.smart_project()
        if was_object_mode:
            bpy.ops.object.mode_set(mode='OBJECT')

    def _select_reference_faces(self, bm) -> None:
        # deselect all
        for face in bm.faces:
            face.select = False 

        # select the reference faces
        for mf in self._matched_faceses:
            ref_fmi = mf.getReferenceFMI()
            face = bm.faces[ref_fmi.face.index]
            face.select = True

    def _stack_matching(self, bm):
        uv_layer = bm.loops.layers.uv.verify()
        for mf in self._matched_faceses:
            ref_fmi = mf.getReferenceFMI()
            for fmi in [fmi for fmi in mf.fmis if fmi is not ref_fmi]:
                face = bm.faces[fmi.face.index]
                ref_face = bm.faces[ref_fmi.face.index]
                for i in range(len(face.loops)):
                    loop = face.loops[(i + fmi.compare_start_index) % len(face.loops)]
                    ref_loop = ref_face.loops[i]
                    loop_uv = loop[uv_layer]
                    ref_loop_uv = ref_loop[uv_layer]
                    loop_uv.uv = ref_loop_uv.uv # assign

    def _uv_map_reference_faces(self, bm):
        bm.faces.ensure_lookup_table() # or else there's an error if we're in object mode
        bm.verts.ensure_lookup_table() # don't need this call?

        self._select_reference_faces(bm)
        self._smart_project()
        self._stack_matching(bm)

    def _populate_fmi_lookup(self, bm) -> None:
        for face in bm.faces:
            fmi = FaceMatchingInfo(face)
            fmi.edge_lengths = [e.calc_length() for e in face.edges]
            self._faces.lookup[face.index] = fmi

    def _find_and_select(self, context):
        ob = context.active_object
        mesh = ob.data
        was_edit_mode = ob.mode == 'EDIT'

        bm = self.bmeshFrom(mesh)
        self._populate_fmi_lookup(bm)
        self._addMatchedFaceSets()
        self._select_reference_faces(bm)
        self._cleanup(bm, mesh)
       
        if was_edit_mode:
            bpy.ops.object.mode_set(mode='EDIT')

    def _stack_on_selected(self, context):
        ob = context.active_object
        mesh = ob.data
        was_edit_mode = ob.mode == 'EDIT'

        bm = self.bmeshFrom(mesh)
        self._populate_fmi_lookup(bm)
        self._sets_from_selected(bm)
        self._onlyMatchWithFaceSets()
        self._stack_matching(bm)
        self._cleanup(bm, mesh)

        if was_edit_mode:
            bpy.ops.object.mode_set(mode='EDIT')

    def _stack(self, context):
        ob = context.active_object
        mesh = ob.data

        was_edit_mode = ob.mode == 'EDIT'
        # we switch to EDIT mode here because it magically
        #  makes a bug go away. The bug is that uv smart project
        #   silently does nothing if we're in OBJ mode (and we don't force switch with this line)
        #  TODO solve this or just require edit mode to run
        bpy.ops.object.mode_set(mode='EDIT') 
       
        # the bmesh approach: https://blender.stackexchange.com/questions/6727/how-to-get-a-list-of-edges-of-current-face-in-bpy
        bm = self.bmeshFrom(mesh)
        self._populate_fmi_lookup(bm)
        self._addMatchedFaceSets()
        
        self._uv_map_reference_faces(bm=bm)

        self._cleanup(bm, mesh)
        if was_edit_mode:
            bpy.ops.object.mode_set(mode='EDIT')

    def _flush(self, bm) -> None:
        bm.select_flush_mode()

    def _cleanup(self, bm, mesh) -> None:
        # apply changes
        self._flush(bm)
        self.apply_bmesh_to_mesh(bm, mesh)
        
        bpy.ops.object.mode_set(mode='OBJECT') # Switching here fends off a bug that crashes blender

        # clean up
        bm.free() 
        bm = None

    @classmethod
    def description(cls, context, p):
        if p.custom_key == 'STACK_ON_SELECTED':
            return "Use selected faces as uv stack targets. Stack other faces onto them if/when they match"
        if p.custom_key == 'FIND_AND_STACK':
            return "Find a set of faces that don't match each other. Smart project those faces. Use them as stack targets for the other faces"
        if p.custom_key == 'FIND_STACK_FACES':
            return "Just find and select the faces that don't match each other. Don't change anything"

    @classmethod
    def poll(cls, context):
        return context.active_object is not None # and context.active_object.mode == 'EDIT' and context.active_object.type == 'MESH'

    def _exec(self, context):
        if self.custom_key == 'STACK_ON_SELECTED':
            self._stack_on_selected(context=context)
            return
        if self.custom_key == 'FIND_AND_STACK':
            self._stack(context=context)
            return
        if self.custom_key == 'FIND_STACK_FACES':
            self._find_and_select(context=context)
            return

    def execute(self, context):
        self._exec(context=context)
        return {'FINISHED'}

class ProtoPanel:
    bl_category = "Mel"
    bl_space_type = "VIEW_3D"
    bl_region_type = "UI"
    bl_options = {"DEFAULT_CLOSED"}

class EXAMPLE_PT_panel_1(ProtoPanel, bpy.types.Panel):
    bl_label = "UV Stacker (experimental)"

    @classmethod
    def poll(cls, context):
        return context.active_object is not None # and context.active_object.mode == 'EDIT'

    def drawProps(self, context):
        layout = self.layout
        layout.prop(context.scene.stack_params, "tolerance")
        layout.prop(context.scene.stack_params, "epsilon")

    def drawButtons(self):
        layout = self.layout
        layout.operator(MelUVStacker.bl_idname, text="find and stack").custom_key = "FIND_AND_STACK"
        layout.operator(MelUVStacker.bl_idname, text="stack onto selected").custom_key = "STACK_ON_SELECTED"
        layout.operator(MelUVStacker.bl_idname, text="find/select stack faces").custom_key = "FIND_STACK_FACES"
        
    def draw(self, context):
        self.drawProps(context=context)
        self.drawButtons()

classes = (EXAMPLE_PT_panel_1, MelUVStacker, MelStackParams)

def register():
    for cls in classes:
        bpy.utils.register_class(cls)
    bpy.types.Scene.stack_params = bpy.props.PointerProperty(type=MelStackParams)

def unregister():
    for cls in classes:
        bpy.utils.unregister_class(cls)
    del bpy.types.Scene.stack_params

if __name__ == "__main__":
    register()