Merge pull request #30 from Yusarina/armature_selection

Armature Selection System

core/common.py

@@ -2,7 +2,7 @@ import bpy
 import numpy as np
 from .dictionaries import bone_names
 
-from typing import List, Optional
+from typing import List, Optional, Tuple
 from bpy.types import Object, ShapeKey, Mesh, Context
 from functools import lru_cache
 
@@ -42,10 +42,10 @@ def has_shapekeys(mesh_obj: Object) -> bool:
 def _get_shape_key_co(shape_key: ShapeKey) -> np.ndarray:
     return np.array([v.co for v in shape_key.data])
     
-def simplify_bonename(n):
+def simplify_bonename(n: str) -> str:
     return n.lower().translate(dict.fromkeys(map(ord, u" _.")))
     
-def get_armature(context, armature_name=None) -> Optional[Object]:
+def get_armature(context: Context, armature_name: Optional[str] = None) -> Optional[Object]:
     if armature_name:
         obj = bpy.data.objects[armature_name]
         if obj.type == "ARMATURE":
@@ -57,3 +57,38 @@ def get_armature(context, armature_name=None) -> Optional[Object]:
         if obj.type == "ARMATURE":
             return obj
     return next((obj for obj in context.view_layer.objects if obj.type == 'ARMATURE'), None)
+
+def get_armatures(self, context: Context) -> List[Tuple[str, str, str]]:
+    return [(obj.name, obj.name, "") for obj in bpy.data.objects if obj.type == 'ARMATURE']
+
+def get_selected_armature(context: Context) -> Optional[Object]:
+    if context.scene.selected_armature:
+        armature = bpy.data.objects.get(context.scene.selected_armature)
+        if is_valid_armature(armature):
+            return armature
+    return None
+
+def set_selected_armature(context: Context, armature: Optional[Object]) -> None:
+    context.scene.selected_armature = armature.name if armature else ""
+
+def is_valid_armature(armature: Object) -> bool:
+    if not armature or armature.type != 'ARMATURE':
+        return False
+    if not armature.data or not armature.data.bones:
+        return False
+    return True
+
+def select_current_armature(context: Context) -> bool:
+    armature = get_selected_armature(context)
+    if armature:
+        bpy.ops.object.select_all(action='DESELECT')
+        armature.select_set(True)
+        context.view_layer.objects.active = armature
+        return True
+    return False
+
+def get_all_meshes(context: Context) -> List[Object]:
+    armature = get_selected_armature(context)
+    if armature and is_valid_armature(armature):
+        return [obj for obj in bpy.data.objects if obj.type == 'MESH' and obj.parent == armature]
+    return []

core/properties.py

@@ -1,6 +1,7 @@
 import bpy
 from ..functions.translations import t, get_languages_list, update_language
 from ..core.addon_preferences import get_preference
+from .common import get_armatures
 
 def register():
     default_language = get_preference("language", 0)
@@ -15,9 +16,18 @@ def register():
     
     bpy.types.Scene.avatar_toolkit_language_changed = bpy.props.BoolProperty(default=False)
 
+    bpy.types.Scene.selected_armature = bpy.props.EnumProperty(
+        items=get_armatures,
+        name="Selected Armature",
+        description="The currently selected armature for Avatar Toolkit operations"
+    )
+
 def unregister():
     if hasattr(bpy.types.Scene, "avatar_toolkit_language"):
         del bpy.types.Scene.avatar_toolkit_language
         
     if hasattr(bpy.types.Scene, "avatar_toolkit_language_changed"):
         del bpy.types.Scene.avatar_toolkit_language_changed
+
+    if hasattr(bpy.types.Scene, "selected_armature"):
+        del bpy.types.Scene.selected_armature

functions/combine_materials.py

@@ -1,8 +1,8 @@
 import bpy
 import re
-from typing import List, Tuple, Optional
+from typing import List, Tuple, Optional, Set
 from bpy.types import Material, Operator, Context, Object
-from ..core.common import clean_material_names
+from ..core.common import clean_material_names, get_selected_armature, is_valid_armature, get_all_meshes
 from ..core.register import register_wrap
 from ..functions.translations import t
 
@@ -65,49 +65,54 @@ class CombineMaterials(Operator):
 
     @classmethod
     def poll(cls, context: Context) -> bool:
-        return context.active_object is not None
+        armature = get_selected_armature(context)
+        return armature is not None and is_valid_armature(armature)
     
-    def execute(self, context: Context) -> set:
-        bpy.ops.object.mode_set(mode='OBJECT')
-        
-        armature: Optional[Object] = next((obj for obj in bpy.data.objects if obj.type == 'ARMATURE'), None)
+    def execute(self, context: Context) -> Set[str]:
+        armature = get_selected_armature(context)
         if not armature:
+            self.report({'WARNING'}, "No armature selected")
             return {'CANCELLED'}
         
-        meshes: List[Object] = [obj for obj in bpy.data.objects if obj.type == 'MESH' and 'Armature' in obj.modifiers and obj.modifiers['Armature'].object == armature]
-        if not meshes:
-            return {'CANCELLED'}
-        
+        context.view_layer.objects.active = armature
         bpy.ops.object.mode_set(mode='OBJECT')
+        
+        meshes = get_all_meshes(context)
+        if not meshes:
+            self.report({'WARNING'}, "No meshes found for the selected armature")
+            return {'CANCELLED'}
+        
         self.consolidate_materials(meshes)
         self.remove_unused_materials()
         self.cleanmatslots()
         self.clean_material_names()
-        bpy.ops.object.mode_set(mode='OBJECT')
-        bpy.context.view_layer.objects.active = armature
         
         return {'FINISHED'}
 
-    def consolidate_materials(self, objects: List[Object]) -> None:
-        mat_mapping: dict = {}
+    def consolidate_materials(self, meshes: List[Object]) -> None:
+        mat_mapping: Dict[str, Material] = {}
         num_combined: int = 0
-        for ob in objects:
-            for slot in ob.material_slots:
+        for mesh in meshes:
+            for slot in mesh.material_slots:
                 mat: Optional[Material] = slot.material
                 if mat:
                     base_name: str = get_base_name(mat.name)
                     
                     if base_name in mat_mapping:
                         base_mat: Material = mat_mapping[base_name]
-                        if materials_match(base_mat, mat):
-                            consolidate_textures(base_mat, mat)
-                            num_combined += 1
-                            slot.material = base_mat
+                        try:
+                            if materials_match(base_mat, mat):
+                                consolidate_textures(base_mat, mat)
+                                num_combined += 1
+                                slot.material = base_mat
+                        except AttributeError:
+                            # Skip this material if there's an attribute mismatch
+                            continue
                     else:
                         mat_mapping[base_name] = mat
         
         report_consolidated(self, num_combined)
-    
+
     def remove_unused_materials(self) -> None:
         for mat in bpy.data.materials:
             if not any(obj for obj in bpy.data.objects if obj.material_slots and mat.name in obj.material_slots):
@@ -125,4 +130,3 @@ class CombineMaterials(Operator):
         for obj in bpy.data.objects:
             if obj.type == 'MESH':
                 clean_material_names(obj)
-

functions/join_meshes.py

@@ -1,8 +1,8 @@
 import bpy
-from typing import List, Optional
+from typing import List, Optional, Set
 from bpy.types import Operator, Context, Object
 from ..core.register import register_wrap
-from ..core.common import fix_uv_coordinates
+from ..core.common import fix_uv_coordinates, get_selected_armature, is_valid_armature, select_current_armature, get_all_meshes
 from ..functions.translations import t
 
 @register_wrap
@@ -14,21 +14,23 @@ class JoinAllMeshes(Operator):
 
     @classmethod
     def poll(cls, context: Context) -> bool:
-        return context.mode == 'OBJECT'
+        armature = get_selected_armature(context)
+        return armature is not None and is_valid_armature(armature)
 
-    def execute(self, context: Context) -> set:
+    def execute(self, context: Context) -> Set[str]:
         self.join_all_meshes(context)
         return {'FINISHED'}
 
     def join_all_meshes(self, context: Context) -> None:
-        if not bpy.data.objects:
-            self.report({'INFO'}, "No objects in the scene")
+        if not select_current_armature(context):
+            self.report({'WARNING'}, "No armature selected")
             return
 
+        armature = get_selected_armature(context)
         bpy.ops.object.mode_set(mode='OBJECT')
         bpy.ops.object.select_all(action='DESELECT')
 
-        meshes: List[Object] = [obj for obj in bpy.data.objects if obj.type == 'MESH']
+        meshes: List[Object] = get_all_meshes(context)
         for mesh in meshes:
             mesh.select_set(True)
 
@@ -43,6 +45,8 @@ class JoinAllMeshes(Operator):
         else:
             self.report({'WARNING'}, "No mesh objects selected")
 
+        context.view_layer.objects.active = armature
+
 @register_wrap
 class JoinSelectedMeshes(Operator):
     bl_idname = "avatar_toolkit.join_selected_meshes"
@@ -52,17 +56,17 @@ class JoinSelectedMeshes(Operator):
 
     @classmethod
     def poll(cls, context: Context) -> bool:
-        return context.mode == 'OBJECT'
+        return context.mode == 'OBJECT' and len([obj for obj in context.selected_objects if obj.type == 'MESH']) > 1
 
-    def execute(self, context: Context) -> set:
+    def execute(self, context: Context) -> Set[str]:
         self.join_selected_meshes(context)
         return {'FINISHED'}
 
     def join_selected_meshes(self, context: Context) -> None:
         selected_objects: List[Object] = [obj for obj in bpy.context.selected_objects if obj.type == 'MESH']
 
-        if not selected_objects:
-            self.report({'WARNING'}, "No mesh objects selected")
+        if len(selected_objects) < 2:
+            self.report({'WARNING'}, "Please select at least two mesh objects")
             return
 
         bpy.ops.object.mode_set(mode='OBJECT')

functions/remove_doubles_safely.py

@@ -5,7 +5,7 @@ import re
 from typing import List, Tuple, Optional, TypedDict
 from bpy.types import Material, Operator, Context, Object
 from ..core.register import register_wrap
-from ..core.common import get_armature
+from ..core.common import get_selected_armature, is_valid_armature, select_current_armature, get_all_meshes
 
 
 class meshEntry(TypedDict):
@@ -23,20 +23,20 @@ class RemoveDoublesSafely(Operator):
 
     @classmethod
     def poll(cls, context: Context) -> bool:
-        return context.mode == 'OBJECT'
+        armature = get_selected_armature(context)
+        return armature is not None and is_valid_armature(armature)
 
     def execute(self, context: Context) -> set:
-        if not bpy.data.objects:
-            self.report({'INFO'}, "No objects in the scene")
-            return
+        if not select_current_armature(context):
+            self.report({'WARNING'}, "No armature selected")
+            return {'CANCELLED'}
 
+        armature = get_selected_armature(context)
         bpy.ops.object.mode_set(mode='OBJECT')
         bpy.ops.object.select_all(action='DESELECT')
-        objects: List[Object] =  get_armature(context).children if get_armature(context) else context.view_layer.objects
+        objects: List[Object] = get_all_meshes(context)
 
-        meshes: List[Object] = [obj for obj in objects if obj.type == 'MESH']  
-        
-        for mesh in meshes:
+        for mesh in objects:
             if mesh.data.name not in [stored_object["mesh"].data.name for stored_object in self.objects_to_do]:
                 mesh_shapekeys = {"mesh":mesh,"shapekeys":[]}
                 mesh_data: bpy.types.Mesh = mesh.data
@@ -45,12 +45,10 @@ class RemoveDoublesSafely(Operator):
                     for shape in mesh_data.shape_keys.key_blocks:
                         mesh_shapekeys["shapekeys"].append(shape.name)
                 self.objects_to_do.append(mesh_shapekeys)
-                
         
         return {'FINISHED'}
         
     def invoke(self, context: Context, event: bpy.types.Event) -> set:
-
         self.execute(context)
         context.window_manager.modal_handler_add(self)
         return {'RUNNING_MODAL'}
@@ -62,7 +60,6 @@ class RemoveDoublesSafely(Operator):
         mesh_data: bpy.types.Mesh = mesh["mesh"].data
         bpy.ops.object.mode_set(mode='EDIT')
         
-        
         bpy.ops.object.mode_set(mode='OBJECT')
         for index, point in enumerate(mesh["mesh"].active_shape_key.points):
             if point.co.xyz != mesh_data.shape_keys.key_blocks[0].points[index].co.xyz:
@@ -70,18 +67,10 @@ class RemoveDoublesSafely(Operator):
                 print("shapekey has a moved vertex at index \""+str(index)+"\", excluding from double merging!")
         bpy.ops.object.mode_set(mode='EDIT')
         
-        
-
-
         bpy.ops.object.mode_set(mode='OBJECT')
         mesh["mesh"].select_set(False)
 
-
     def modal(self, context: Context, event: bpy.types.Event) -> set:
-        
-        
-        
-
         if len(self.objects_to_do) > 0:
             mesh = self.objects_to_do[0]
             mesh_data: bpy.types.Mesh = mesh["mesh"].data
@@ -131,10 +120,3 @@ class RemoveDoublesSafely(Operator):
             return {'FINISHED'}
         
         return {'RUNNING_MODAL'}
-                
-                
-                
-
-
-
-

functions/resonite_functions.py

@@ -4,7 +4,7 @@ from typing import List, Optional
 import re
 from bpy.types import Operator, Context, Object
 from ..core.dictionaries import bone_names
-from ..core.common import get_armature, simplify_bonename
+from ..core.common import get_selected_armature, simplify_bonename, is_valid_armature
 from ..functions.translations import t
 
 @register_wrap
@@ -16,12 +16,14 @@ class ConvertToResonite(Operator):
 
     @classmethod
     def poll(cls, context: Context) -> bool:
-        if not get_armature(context):
-            return False
-        return True
+        armature = get_selected_armature(context)
+        return armature is not None and is_valid_armature(armature)
         
     def execute(self, context: Context) -> set:
-        armature = get_armature(context)
+        armature = get_selected_armature(context)
+        if not armature:
+            self.report({'WARNING'}, "No armature selected")
+            return {'CANCELLED'}
 
         translate_bone_fails = 0
         untranslated_bones = set()
@@ -89,8 +91,6 @@ class ConvertToResonite(Operator):
             'thumb_3_r': "thumb3.R"
         }
 
-                
-
         context.view_layer.objects.active = armature
         bpy.ops.object.mode_set(mode='EDIT')
 
@@ -111,5 +111,4 @@ class ConvertToResonite(Operator):
         else:
             self.report({'INFO'}, "Successfully translated all bones to humanoid names")
 
-
-        return {'FINISHED'}
+        return {'FINISHED'}

ui/optimization.py

@@ -2,6 +2,7 @@ import bpy
 from ..core.register import register_wrap
 from .panel import AvatarToolkitPanel
 from ..functions.translations import t
+from ..core.common import get_selected_armature
 
 @register_wrap
 class AvatarToolkitOptimizationPanel(bpy.types.Panel):
@@ -14,20 +15,21 @@ class AvatarToolkitOptimizationPanel(bpy.types.Panel):
 
     def draw(self, context):
         layout = self.layout
-        layout.label(text=t("Optimization.options.label"))
+        armature = get_selected_armature(context)
         
-        row = layout.row()
-        row.scale_y = 1.2 
-        row.operator("avatar_toolkit.combine_materials", text=t("Optimization.combine_materials.label"))
-        
-
-        layout.separator(factor=0.5)
-        
-        row = layout.row(align=True)
-        row.scale_y = 1.2 
-        row.operator("avatar_toolkit.join_all_meshes", text=t("Optimization.join_all_meshes.label"))
-        row.operator("avatar_toolkit.join_selected_meshes", text=t("Optimization.join_selected_meshes.label"))
-        row.operator("avatar_toolkit.remove_doubles_safely", text="Remove Doubles Safely")
-
-        # Add optimization options here
-
+        if armature:
+            layout.label(text=t("Optimization.options.label"))
+            
+            row = layout.row()
+            row.scale_y = 1.2 
+            row.operator("avatar_toolkit.combine_materials", text=t("Optimization.combine_materials.label"))
+            
+            layout.separator(factor=0.5)
+            
+            row = layout.row(align=True)
+            row.scale_y = 1.2 
+            row.operator("avatar_toolkit.join_all_meshes", text=t("Optimization.join_all_meshes.label"))
+            row.operator("avatar_toolkit.join_selected_meshes", text=t("Optimization.join_selected_meshes.label"))
+            row.operator("avatar_toolkit.remove_doubles_safely", text="Remove Doubles Safely")
+        else:
+            layout.label(text="Please select an armature in Quick Access")

ui/quick_access.py

@@ -7,6 +7,7 @@ from ..functions.translations import t
 from ..core.import_pmx import import_pmx
 from ..core.import_pmd import import_pmd
 from ..core.importer import import_fbx
+from ..core.common import get_selected_armature, set_selected_armature
 
 @register_wrap
 class AvatarToolkitQuickAccessPanel(bpy.types.Panel):
@@ -21,6 +22,8 @@ class AvatarToolkitQuickAccessPanel(bpy.types.Panel):
         layout = self.layout
         layout.label(text=t("Quick_Access.options"))
 
+        layout.prop(context.scene, "selected_armature", text="Select Armature")
+
         row = layout.row()
         row.label(text=t("Quick_Access.import_export.label"), icon='IMPORT')
         
@@ -129,5 +132,3 @@ class AVATAR_TOOLKIT_OT_export_fbx(bpy.types.Operator):
     def execute(self, context):
         bpy.ops.export_scene.fbx('INVOKE_DEFAULT')
         return {'FINISHED'}
-
-

ui/tools.py

@@ -3,6 +3,7 @@ from ..core.register import register_wrap
 from .panel import AvatarToolkitPanel
 from bpy.types import Context
 from ..functions.translations import t
+from ..core.common import get_selected_armature
 
 @register_wrap
 class AvatarToolkitToolsPanel(bpy.types.Panel):
@@ -15,11 +16,16 @@ class AvatarToolkitToolsPanel(bpy.types.Panel):
 
     def draw(self, context: Context):
         layout = self.layout
-        layout.label(text=t("Tools.tools_title.label"))
-        layout.separator(factor=0.5)
+        armature = get_selected_armature(context)
+        
+        if armature:
+            layout.label(text=t("Tools.tools_title.label"))
+            layout.separator(factor=0.5)
 
-        row = layout.row(align=True)
-        row.scale_y = 1.5  
-        row.operator("avatar_toolkit.convert_to_resonite", text=t("Tools.convert_to_resonite.label"))
-        row = layout.row(align=True)
-        row.operator("avatar_toolkit.remove_doubles_safely", text="Remove Doubles Safely")
+            row = layout.row(align=True)
+            row.scale_y = 1.5  
+            row.operator("avatar_toolkit.convert_to_resonite", text=t("Tools.convert_to_resonite.label"))
+            row = layout.row(align=True)
+            row.operator("avatar_toolkit.remove_doubles_safely", text="Remove Doubles Safely")
+        else:
+            layout.label(text="Please select an armature in Quick Access")