backend/scripts/blender/rig_quadruped_blender.py

import sys
import os

try:
    import bpy
    import mathutils
except ImportError:
    print("[Error] This script must be run inside Blender's python environment.")
    sys.exit(1)

def strip_gltf_extensions(glb_path):
    try:
        with open(glb_path, "rb") as f:
            data = f.read()
        if len(data) < 20 or data[:4] != b'glTF':
            return
        json_len = int.from_bytes(data[12:16], byteorder='little')
        json_bytes = data[20:20+json_len]
        import json
        gltf_json = json.loads(json_bytes.decode('utf-8', errors='ignore'))
        
        modified = False
        for key in ['extensionsRequired', 'extensionsUsed']:
            if key in gltf_json and 'EXT_texture_webp' in gltf_json[key]:
                gltf_json[key].remove('EXT_texture_webp')
                modified = True
                
        if modified:
            new_bytes = json.dumps(gltf_json).encode('utf-8')
            if len(new_bytes) <= len(json_bytes):
                new_bytes = new_bytes.ljust(len(json_bytes), b' ')
                new_data = data[:20] + new_bytes + data[20+len(json_bytes):]
                with open(glb_path, "wb") as f:
                    f.write(new_data)
                print(f"[Blender Quadruped] Stripped EXT_texture_webp extension requirement from GLB.")
    except Exception as e:
        print(f"[Blender Quadruped] Extension strip note: {e}")

def run_blender_rigging(input_path, output_path):
    print(f"[Blender Rigging] Loading model: {input_path}")
    
    # 1. Reset scene
    bpy.ops.wm.read_factory_settings(use_empty=True)
    
    # 2. Import GLB
    strip_gltf_extensions(input_path)
    bpy.ops.import_scene.gltf(filepath=input_path)
    
    # 3. Find all mesh objects and join them into a single character mesh
    mesh_objs = [obj for obj in bpy.context.scene.objects if obj.type == 'MESH']
    if not mesh_objs:
        print("[Error] No mesh found in imported GLTF.")
        sys.exit(1)
        
    print(f"[Blender Rigging] Found {len(mesh_objs)} meshes. Joining them...")
    bpy.ops.object.select_all(action='DESELECT')
    for obj in mesh_objs:
        obj.select_set(True)
    bpy.context.view_layer.objects.active = mesh_objs[0]
    bpy.ops.object.join()
    
    mesh_obj = bpy.context.active_object
    mesh_obj.name = "CharacterMesh"
    print(f"[Blender Rigging] Joined mesh: {mesh_obj.name}")
    
    # Ensure transforms are applied so bounding box calculations are correct
    bpy.ops.object.select_all(action='DESELECT')
    mesh_obj.select_set(True)
    bpy.context.view_layer.objects.active = mesh_obj
    bpy.ops.object.transform_apply(location=True, rotation=True, scale=True)
    
    # 4. Calculate bounding box
    bbox = [mesh_obj.matrix_world @ mathutils.Vector(corner) for corner in mesh_obj.bound_box]
    min_x = min([v.x for v in bbox])
    max_x = max([v.x for v in bbox])
    min_y = min([v.y for v in bbox])
    max_y = max([v.y for v in bbox])
    min_z = min([v.z for v in bbox])
    max_z = max([v.z for v in bbox])
    
    width = max_x - min_x
    depth = max_y - min_y
    height = max_z - min_z
    center_x = (min_x + max_x) / 2.0
    center_y = (min_y + max_y) / 2.0
    
    print(f"[Blender Rigging] Mesh Bounds: Height={height:.4f}, Width={width:.4f}, Depth={depth:.4f}")
    print(f"[Blender Rigging] Mesh Center: X={center_x:.4f}, Y={center_y:.4f}, Min Z={min_z:.4f}")
    
    # 5. Define relative quadruped bone joints
    # positive Y in Blender goes back, negative Y goes front
    z_spine = min_z + height * 0.65
    z_head = min_z + height * 0.85
    z_neck = min_z + height * 0.75
    z_rear_hip = min_z + height * 0.60
    z_rear_knee = min_z + height * 0.35
    z_rear_ankle = min_z + height * 0.12
    z_rear_foot = min_z + height * 0.02
    
    z_front_shoulder = min_z + height * 0.62
    z_front_elbow = min_z + height * 0.35
    z_front_wrist = min_z + height * 0.12
    z_front_foot = min_z + height * 0.02
    
    y_rear = center_y + depth * 0.28
    y_front = center_y - depth * 0.25
    y_head = center_y - depth * 0.42
    y_neck = center_y - depth * 0.30
    y_tail_base = center_y + depth * 0.32
    y_tail_tip = center_y + depth * 0.55
    
    x_offset_rear = width * 0.22
    x_offset_front = width * 0.22
    
    x_rear_l = center_x - x_offset_rear
    x_rear_r = center_x + x_offset_rear
    x_front_l = center_x - x_offset_front
    x_front_r = center_x + x_offset_front
    
    joints = {
        # Spine & Head
        "Pelvis": ((center_x, y_rear, z_spine), (center_x, center_y, z_spine)),
        "Spine": ((center_x, center_y, z_spine), (center_x, y_front, z_spine)),
        "Neck": ((center_x, y_front, z_spine), (center_x, y_neck, z_neck)),
        "Head": ((center_x, y_neck, z_neck), (center_x, y_head, z_head)),
        "Tail": ((center_x, y_tail_base, z_spine), (center_x, y_tail_tip, z_spine - height * 0.1)),
        
        # Left Rear Leg
        "LeftUpLegRear": ((x_rear_l, y_rear, z_rear_hip), (x_rear_l, y_rear, z_rear_knee)),
        "LeftLegRear": ((x_rear_l, y_rear, z_rear_knee), (x_rear_l, y_rear, z_rear_ankle)),
        "LeftFootRear": ((x_rear_l, y_rear, z_rear_ankle), (x_rear_l, y_rear - depth * 0.05, z_rear_foot)),
        
        # Right Rear Leg
        "RightUpLegRear": ((x_rear_r, y_rear, z_rear_hip), (x_rear_r, y_rear, z_rear_knee)),
        "RightLegRear": ((x_rear_r, y_rear, z_rear_knee), (x_rear_r, y_rear, z_rear_ankle)),
        "RightFootRear": ((x_rear_r, y_rear, z_rear_ankle), (x_rear_r, y_rear - depth * 0.05, z_rear_foot)),
        
        # Left Front Leg
        "LeftUpLegFront": ((x_front_l, y_front, z_front_shoulder), (x_front_l, y_front, z_front_elbow)),
        "LeftLegFront": ((x_front_l, y_front, z_front_elbow), (x_front_l, y_front, z_front_wrist)),
        "LeftFootFront": ((x_front_l, y_front, z_front_wrist), (x_front_l, y_front - depth * 0.05, z_front_foot)),
        
        # Right Front Leg
        "RightUpLegFront": ((x_front_r, y_front, z_front_shoulder), (x_front_r, y_front, z_front_elbow)),
        "RightLegFront": ((x_front_r, y_front, z_front_elbow), (x_front_r, y_front, z_front_wrist)),
        "RightFootFront": ((x_front_r, y_front, z_front_wrist), (x_front_r, y_front - depth * 0.05, z_front_foot))
    }
    
    # 6. Create Armature
    print("[Blender Rigging] Creating Armature...")
    arm_data = bpy.data.armatures.new(name="QuadrupedArmature")
    rig_obj = bpy.data.objects.new(name="QuadrupedRig", object_data=arm_data)
    bpy.context.scene.collection.objects.link(rig_obj)
    
    bpy.context.view_layer.objects.active = rig_obj
    bpy.ops.object.mode_set(mode='EDIT')
    
    edit_bones = arm_data.edit_bones
    bone_objects = {}
    for name, (head, tail) in joints.items():
        bone = edit_bones.new(name)
        bone.head = head
        bone.tail = tail
        bone_objects[name] = bone
        
    # Setup hierarchy
    bone_objects["Spine"].parent = bone_objects["Pelvis"]
    bone_objects["Neck"].parent = bone_objects["Spine"]
    bone_objects["Head"].parent = bone_objects["Neck"]
    bone_objects["Tail"].parent = bone_objects["Pelvis"]
    
    # Rear legs hierarchy
    bone_objects["LeftUpLegRear"].parent = bone_objects["Pelvis"]
    bone_objects["LeftLegRear"].parent = bone_objects["LeftUpLegRear"]
    bone_objects["LeftFootRear"].parent = bone_objects["LeftLegRear"]
    
    bone_objects["RightUpLegRear"].parent = bone_objects["Pelvis"]
    bone_objects["RightLegRear"].parent = bone_objects["RightUpLegRear"]
    bone_objects["RightFootRear"].parent = bone_objects["RightLegRear"]
    
    # Front legs hierarchy
    bone_objects["LeftUpLegFront"].parent = bone_objects["Spine"]
    bone_objects["LeftLegFront"].parent = bone_objects["LeftUpLegFront"]
    bone_objects["LeftFootFront"].parent = bone_objects["LeftLegFront"]
    
    bone_objects["RightUpLegFront"].parent = bone_objects["Spine"]
    bone_objects["RightLegFront"].parent = bone_objects["RightUpLegFront"]
    bone_objects["RightFootFront"].parent = bone_objects["RightLegFront"]
    
    bpy.ops.object.mode_set(mode='OBJECT')
    
    # 7. Parent mesh to Armature using a Voxel Proxy for 100% reliable weighting
    print("[Blender Rigging] Creating watertight Voxel Proxy mesh for auto-weighting calculation...")
    bpy.ops.object.select_all(action='DESELECT')
    mesh_obj.select_set(True)
    bpy.context.view_layer.objects.active = mesh_obj
    bpy.ops.object.duplicate(linked=False)
    proxy_obj = bpy.context.active_object
    proxy_obj.name = "VoxelProxyMesh"
    
    bbox_size = max(proxy_obj.dimensions)
    voxel_size = max(0.003, bbox_size / 150.0)
    print(f"[Blender Rigging] Remeshing proxy with voxel size: {voxel_size:.4f}")
    
    try:
        proxy_obj.data.remesh_voxel_size = voxel_size
        bpy.ops.object.voxel_remesh()
        print("[Blender Rigging] ✓ Voxel Remesh completed on proxy.")
    except Exception as re_err:
        print(f"[Blender Rigging] Voxel Remesh failed on proxy: {re_err}")
        
    bpy.ops.object.select_all(action='DESELECT')
    proxy_obj.select_set(True)
    rig_obj.select_set(True)
    bpy.context.view_layer.objects.active = rig_obj
    
    proxy_rig_success = False
    try:
        bpy.ops.object.parent_set(type='ARMATURE_AUTO')
        print("[Blender Rigging] ✓ Parented Voxel Proxy with automatic weights successfully!")
        proxy_rig_success = True
    except Exception as parent_err:
        print(f"[Blender Rigging] Error: Auto weighting failed even on proxy: {parent_err}")
        
    if proxy_rig_success:
        print("[Blender Rigging] Transferring skin weights from Voxel Proxy to original detailed mesh...")
        bpy.ops.object.select_all(action='DESELECT')
        mesh_obj.select_set(True)
        rig_obj.select_set(True)
        bpy.context.view_layer.objects.active = rig_obj
        bpy.ops.object.parent_set(type='ARMATURE_NAME')
        
        bpy.ops.object.select_all(action='DESELECT')
        mesh_obj.select_set(True)
        bpy.context.view_layer.objects.active = mesh_obj
        
        dt_mod = mesh_obj.modifiers.new(name="WeightTransfer", type='DATA_TRANSFER')
        dt_mod.object = proxy_obj
        dt_mod.use_vert_data = True
        dt_mod.data_types_verts = {'VGROUP_WEIGHTS'}
        dt_mod.vert_mapping = 'POLYINTERP_NEAREST'
        
        bpy.ops.object.datalayout_transfer(modifier="WeightTransfer")
        bpy.ops.object.modifier_apply(modifier="WeightTransfer")
        print("[Blender Rigging] ✓ Weights transferred successfully.")
        
        bpy.ops.object.select_all(action='DESELECT')
        proxy_obj.select_set(True)
        bpy.ops.object.delete()
        print("[Blender Rigging] Deleted temporary Voxel Proxy.")
    else:
        print("[Blender Rigging] Falling back to default empty weights parenting on original mesh...")
        bpy.ops.object.select_all(action='DESELECT')
        mesh_obj.select_set(True)
        rig_obj.select_set(True)
        bpy.context.view_layer.objects.active = rig_obj
        bpy.ops.object.parent_set(type='ARMATURE')
            
    # 8. Export rigged model to FBX
    print(f"[Blender Rigging] Exporting rigged model to FBX: {output_path}")
    bpy.ops.object.select_all(action='DESELECT')
    mesh_obj.select_set(True)
    rig_obj.select_set(True)
    
    try:
        bpy.ops.export_scene.fbx(
            filepath=output_path,
            use_selection=True,
            path_mode='COPY',
            embed_textures=True,
            add_leaf_bones=False
        )
        print("[Blender Rigging] ✓ Rigged FBX exported successfully.")
    except Exception as fbx_err:
        print(f"[Blender Rigging] Error: FBX export failed: {fbx_err}")
        sys.exit(1)

    # Export rigged GLB for Web visualization
    glb_output_path = output_path.replace('.fbx', '.glb')
    print(f"[Blender Rigging] Exporting rigged model to GLB: {glb_output_path}")
    try:
        bpy.ops.export_scene.gltf(
            filepath=glb_output_path,
            export_format='GLB',
            use_selection=True
        )
        print("[Blender Rigging] ✓ Rigged GLB exported successfully.")
    except Exception as glb_err:
        print(f"[Blender Rigging] Warning: GLB export failed: {glb_err}")

if __name__ == "__main__":
    try:
        args_idx = sys.argv.index("--")
        args = sys.argv[args_idx + 1:]
    except ValueError:
        args = []
        
    if len(args) < 2:
        print("[Error] Usage: blender --background --python rig_quadruped_blender.py -- <input_glb> <output_fbx>")
        sys.exit(1)
        
    run_blender_rigging(args[0], args[1])