blender-mixamoo-openpose-attach/mixamo_openpose.py

bl_info = {
    "name": "OpenPose Mixamo Connect",
    "author": "Your Name",
    "version": (1, 0, 0),
    "blender": (3, 0, 0),
    "location": "View3D > Sidebar > OpenPose",
    "description": "Connect OpenPose rig to Mixamo armature with bone mapping",
    "warning": "",
    "doc_url": "",
    "category": "Rigging",
}

import bpy
import os
from bpy.types import Panel, Operator, PropertyGroup
from bpy.props import PointerProperty, StringProperty


# Bone mapping dictionary
BONE_MAP = {
    'Head': 'mixamorig:Head',
    'Foot_R': 'mixamorig:RightFoot',
    'Leg_R': 'mixamorig:RightLeg',  
    'Foot_L': 'mixamorig:LeftFoot',
    'Leg_L': 'mixamorig:LeftLeg',
    'Pelvis_R': 'mixamorig:RightUpLeg',
    'Pelvis_L': 'mixamorig:LeftUpLeg',
    'UpperArm_R': 'mixamorig:RightArm',
    'UpperArm_L': 'mixamorig:LeftArm',
    'ForeArm_R': 'mixamorig:RightForeArm',
    'ForeArm_L': 'mixamorig:LeftForeArm',
    'Hand_R': 'mixamorig:RightHand',
    'Hand_L': 'mixamorig:LeftHand',
    'Bone.007': 'mixamorig:Spine2',
}


class OPENPOSE_OT_add_rig(Operator):
    """Add OpenPose rig to the scene from bundled blend file"""
    bl_idname = "openpose.add_rig"
    bl_label = "Add OpenPose Rig"
    bl_options = {'REGISTER', 'UNDO'}
    
    def execute(self, context):
        # Get the addon directory
        addon_dir = os.path.dirname(os.path.realpath(__file__))
        blend_file = os.path.join(addon_dir, "openpose.blend")
        blend_file = "/home/noti/dev/blender_mixamo_openpose_attach/openpose.blend"
        
        # Check if the blend file exists
        if not os.path.exists(blend_file):
            self.report({'ERROR'}, f"OpenPose rig file not found: {blend_file}")
            return {'CANCELLED'}
        
        # Try to append the OpenPose rig
        try:
            # Get all existing objects before appending
            existing_objects = set(bpy.data.objects)
            
            # Append all objects from the blend file
            with bpy.data.libraries.load(blend_file, link=False) as (data_from, data_to):
                data_to.objects = data_from.objects
            
            # Add the appended objects to the current scene
            new_objects = []
            for obj in data_to.objects:
                if obj is not None and obj not in existing_objects:
                    context.collection.objects.link(obj)
                    new_objects.append(obj)
            
            # Select the newly added objects
            bpy.ops.object.select_all(action='DESELECT')
            for obj in new_objects:
                obj.select_set(True)
                if obj.type == 'ARMATURE':
                    context.view_layer.objects.active = obj
            
            self.report({'INFO'}, f"Added OpenPose rig: {len(new_objects)} object(s)")
            return {'FINISHED'}
            
        except Exception as e:
            self.report({'ERROR'}, f"Failed to load OpenPose rig: {str(e)}")
            return {'CANCELLED'}


class OPENPOSE_OT_connect_rigs(Operator):
    """Connect OpenPose rig to Mixamo armature using bone constraints"""
    bl_idname = "openpose.connect_rigs"
    bl_label = "Connect Rigs"
    bl_options = {'REGISTER', 'UNDO'}
    
    @classmethod
    def poll(cls, context):
        props = context.scene.openpose_props
        return (props.source_armature is not None and 
                props.target_armature is not None and
                props.source_armature.type == 'ARMATURE' and
                props.target_armature.type == 'ARMATURE')
    
    def execute(self, context):
        props = context.scene.openpose_props
        source_armature = props.source_armature
        target_armature = props.target_armature
        
        if source_armature == target_armature:
            self.report({'ERROR'}, "Source and target armatures must be different")
            return {'CANCELLED'}
        
        # Statistics
        mapped_count = 0
        missing_source = []
        missing_target = []
        
        # Loop through the bone map
        for source_bone_name, target_bone_name in BONE_MAP.items():
            # Get the pose bones
            source_bone = source_armature.pose.bones.get(source_bone_name)
            target_bone = target_armature.pose.bones.get(target_bone_name)
            
            if not source_bone:
                missing_source.append(source_bone_name)
                continue
                
            if not target_bone:
                missing_target.append(target_bone_name)
                continue
            
            # Clear existing Copy Location constraints from source bone
            for constraint in list(source_bone.constraints):
                if constraint.type == 'COPY_LOCATION':
                    source_bone.constraints.remove(constraint)
            
            # Add Copy Location constraint
            cl_constraint = source_bone.constraints.new('COPY_LOCATION')
            cl_constraint.target = target_armature
            cl_constraint.subtarget = target_bone_name
            cl_constraint.use_offset = False
            
            mapped_count += 1
        
        # Update scene
        context.view_layer.update()
        
        # Report results
        report_msg = f"Connected {mapped_count} bones"
        if missing_source:
            report_msg += f" | Missing in source: {len(missing_source)}"
        if missing_target:
            report_msg += f" | Missing in target: {len(missing_target)}"
        
        self.report({'INFO'}, report_msg)
        
        if missing_source or missing_target:
            print(f"\n=== OpenPose Mixamo Connect Report ===")
            if missing_source:
                print(f"Missing source bones: {', '.join(missing_source)}")
            if missing_target:
                print(f"Missing target bones: {', '.join(missing_target)}")
        
        return {'FINISHED'}


class OPENPOSE_OT_clear_constraints(Operator):
    """Clear all Copy Location constraints from the source armature"""
    bl_idname = "openpose.clear_constraints"
    bl_label = "Clear Constraints"
    bl_options = {'REGISTER', 'UNDO'}
    
    @classmethod
    def poll(cls, context):
        props = context.scene.openpose_props
        return props.source_armature is not None and props.source_armature.type == 'ARMATURE'
    
    def execute(self, context):
        props = context.scene.openpose_props
        source_armature = props.source_armature
        
        cleared_count = 0
        for bone in source_armature.pose.bones:
            for constraint in list(bone.constraints):
                if constraint.type == 'COPY_LOCATION':
                    bone.constraints.remove(constraint)
                    cleared_count += 1
        
        context.view_layer.update()
        self.report({'INFO'}, f"Cleared {cleared_count} constraint(s)")
        return {'FINISHED'}


class OPENPOSE_PT_main_panel(Panel):
    """Main panel for OpenPose Mixamo Connect addon"""
    bl_label = "OpenPose Mixamo Connect"
    bl_idname = "OPENPOSE_PT_main_panel"
    bl_space_type = 'VIEW_3D'
    bl_region_type = 'UI'
    bl_category = 'OpenPose'
    
    def draw(self, context):
        layout = self.layout
        props = context.scene.openpose_props
        
        # Add OpenPose Rig section
        box = layout.box()
        box.label(text="1. Add OpenPose Rig", icon='ARMATURE_DATA')
        box.operator("openpose.add_rig", icon='ADD')
        
        layout.separator()
        
        # Select Armatures section
        box = layout.box()
        box.label(text="2. Select Armatures", icon='OBJECT_DATA')
        
        col = box.column(align=True)
        col.label(text="Source (OpenPose):")
        col.prop(props, "source_armature", text="")
        
        col.separator()
        
        col.label(text="Target (Mixamo):")
        col.prop(props, "target_armature", text="")
        
        layout.separator()
        
        # Connect Rigs section
        box = layout.box()
        box.label(text="3. Connect Rigs", icon='LINKED')
        
        col = box.column(align=True)
        col.operator("openpose.connect_rigs", icon='CON_LOCLIKE')
        col.operator("openpose.clear_constraints", icon='X')
        
        # Info section
        if props.source_armature and props.target_armature:
            layout.separator()
            info_box = layout.box()
            info_box.label(text=f"Bone mappings: {len(BONE_MAP)}", icon='INFO')


class OpenPoseProperties(PropertyGroup):
    """Property group to store addon properties"""
    source_armature: PointerProperty(
        name="Source Armature",
        description="The OpenPose armature (will receive constraints)",
        type=bpy.types.Object,
        poll=lambda self, obj: obj.type == 'ARMATURE'
    )
    
    target_armature: PointerProperty(
        name="Target Armature",
        description="The Mixamo armature (constraint target)",
        type=bpy.types.Object,
        poll=lambda self, obj: obj.type == 'ARMATURE'
    )


# Registration
classes = (
    OpenPoseProperties,
    OPENPOSE_OT_add_rig,
    OPENPOSE_OT_connect_rigs,
    OPENPOSE_OT_clear_constraints,
    OPENPOSE_PT_main_panel,
)


def register():
    for cls in classes:
        bpy.utils.register_class(cls)
    
    bpy.types.Scene.openpose_props = PointerProperty(type=OpenPoseProperties)


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


if __name__ == "__main__":
    register()