settled-reach/spikes/quaternius-aesthetic/scripts/blender/separate_head.py

"""
SUPERSEDED: Head-only separation from a Quaternius body mesh.

This script splits the body into head vs body based on Head bone weights.
It has been superseded by segment_body.py, which performs full 15-region
segmentation (including the head as one of those regions).

Use segment_body.py for all new work. This file is preserved as documentation.

Run via:
    reach blender run \\
        spikes/quaternius-aesthetic/scripts/blender/separate_head.py \\
        -- <input.gltf> <body_output.glb> <head_output.glb>
"""

import bpy
import bmesh
import sys
import os

argv = sys.argv
argv = argv[argv.index("--") + 1:] if "--" in argv else []

if len(argv) < 3:
    print("Usage: -- <input.gltf> <body_output.glb> <head_output.glb>")
    sys.exit(1)

INPUT_PATH = argv[0]
BODY_OUTPUT = argv[1]
HEAD_OUTPUT = argv[2]

HEAD_BONE = "Head"

print("=== Separate head from body ===")
print(f"  Input:  {INPUT_PATH}")
print(f"  Body:   {BODY_OUTPUT}")
print(f"  Head:   {HEAD_OUTPUT}")


def clear_scene():
    bpy.ops.object.select_all(action='SELECT')
    bpy.ops.object.delete()
    for c in list(bpy.data.collections):
        bpy.data.collections.remove(c)


def get_dominant_bone(mesh_obj, vert_index):
    """Get the name of the bone with highest weight for a vertex."""
    best_weight = 0.0
    best_group = ""
    for g in mesh_obj.data.vertices[vert_index].groups:
        if g.weight > best_weight:
            best_weight = g.weight
            vg = mesh_obj.vertex_groups[g.group]
            best_group = vg.name
    return best_group


def separate_by_bone(mesh_obj, bone_name):
    """
    Separate vertices weighted to bone_name into a new object.
    Returns (body_obj, head_obj).
    """
    # Find which vertices belong to the head
    head_verts = set()
    for v in mesh_obj.data.vertices:
        dominant = get_dominant_bone(mesh_obj, v.index)
        if dominant == bone_name:
            head_verts.add(v.index)

    print(f"  Total verts: {len(mesh_obj.data.vertices)}")
    print(f"  Head verts: {len(head_verts)}")
    print(f"  Body verts: {len(mesh_obj.data.vertices) - len(head_verts)}")

    if not head_verts:
        print("  WARNING: no vertices found for head bone")
        return mesh_obj, None

    # Duplicate the mesh for the head version
    bpy.ops.object.select_all(action='DESELECT')
    mesh_obj.select_set(True)
    bpy.context.view_layer.objects.active = mesh_obj
    bpy.ops.object.duplicate()
    head_obj = bpy.context.active_object
    head_obj.name = "Head_Mesh"

    # Edit the BODY copy — remove head vertices
    bpy.ops.object.select_all(action='DESELECT')
    mesh_obj.select_set(True)
    bpy.context.view_layer.objects.active = mesh_obj
    bpy.ops.object.mode_set(mode='EDIT')
    bpy.ops.mesh.select_all(action='DESELECT')
    bpy.ops.object.mode_set(mode='OBJECT')

    for v in mesh_obj.data.vertices:
        v.select = v.index in head_verts

    bpy.ops.object.mode_set(mode='EDIT')
    bpy.ops.mesh.delete(type='VERT')
    bpy.ops.object.mode_set(mode='OBJECT')

    # Edit the HEAD copy — remove body vertices
    bpy.ops.object.select_all(action='DESELECT')
    head_obj.select_set(True)
    bpy.context.view_layer.objects.active = head_obj
    bpy.ops.object.mode_set(mode='EDIT')
    bpy.ops.mesh.select_all(action='DESELECT')
    bpy.ops.object.mode_set(mode='OBJECT')

    # Head verts in the duplicate have the same indices
    body_verts = set(range(len(head_obj.data.vertices))) - head_verts
    for v in head_obj.data.vertices:
        v.select = v.index in body_verts

    bpy.ops.object.mode_set(mode='EDIT')
    bpy.ops.mesh.delete(type='VERT')
    bpy.ops.object.mode_set(mode='OBJECT')

    print(f"  Body mesh after split: {len(mesh_obj.data.vertices)} verts")
    print(f"  Head mesh after split: {len(head_obj.data.vertices)} verts")

    return mesh_obj, head_obj


def export_glb(objects, path):
    """Export selected objects as GLB."""
    bpy.ops.object.select_all(action='DESELECT')
    for obj in objects:
        obj.select_set(True)
        obj.hide_set(False)
    bpy.context.view_layer.objects.active = objects[0]
    os.makedirs(os.path.dirname(path) or ".", exist_ok=True)
    bpy.ops.export_scene.gltf(
        filepath=path,
        export_format='GLB',
        use_selection=True,
        export_apply=False,
        export_animations=False,
        export_skins=True,
    )
    print(f"  Exported: {path} ({os.path.getsize(path) // 1024} KB)")


# --- Main ---

clear_scene()

print("\nStep 1: Import...")
bpy.ops.import_scene.gltf(filepath=INPUT_PATH)
all_objects = list(bpy.data.objects)

armature = None
meshes = []
other = []
for obj in all_objects:
    if obj.type == 'ARMATURE':
        armature = obj
    elif obj.type == 'MESH':
        meshes.append(obj)
    else:
        other.append(obj)

if not armature:
    print("ERROR: no armature found")
    sys.exit(1)

print(f"  Armature: {armature.name}")
for m in meshes:
    print(f"  Mesh: {m.name} ({len(m.data.vertices)} verts)")

# Find the main body mesh (largest)
body_mesh = max(meshes, key=lambda m: len(m.data.vertices))
small_meshes = [m for m in meshes if m != body_mesh]  # eyes, eyebrows

print(f"\nStep 2: Separate head from {body_mesh.name}...")
body_only, head_only = separate_by_bone(body_mesh, HEAD_BONE)

if head_only is None:
    print("ERROR: head separation failed")
    sys.exit(1)

# Step 3: Export body (armature + body mesh + eyes/eyebrows — no head)
print("\nStep 3: Export body (without head)...")
export_objects = [armature, body_only] + small_meshes
# Hide head for body export
head_only.hide_set(True)
export_glb(export_objects, BODY_OUTPUT)

# Step 4: Export head as standalone mesh (no armature, centered at origin)
print("\nStep 4: Export head...")
body_only.hide_set(True)
for m in small_meshes:
    m.hide_set(True)
head_only.hide_set(False)

# Remove armature modifier and unparent from skeleton
bpy.context.view_layer.objects.active = head_only
head_only.select_set(True)
for mod in list(head_only.modifiers):
    if mod.type == 'ARMATURE':
        bpy.ops.object.modifier_apply(modifier=mod.name)
bpy.ops.object.parent_clear(type='CLEAR_KEEP_TRANSFORM')

# Find the Head bone position to use as the new origin
head_bone = armature.data.bones.get(HEAD_BONE)
if head_bone:
    head_bone_pos = armature.matrix_world @ head_bone.head_local
    print(f"  Head bone position: {head_bone_pos}")
    # Set origin to head bone position, then move mesh to world origin
    bpy.context.scene.cursor.location = head_bone_pos
    bpy.ops.object.origin_set(type='ORIGIN_CURSOR')
    head_only.location = (0, 0, 0)

bpy.ops.object.select_all(action='DESELECT')
head_only.select_set(True)
bpy.context.view_layer.objects.active = head_only
bpy.ops.export_scene.gltf(
    filepath=HEAD_OUTPUT,
    export_format='GLB',
    use_selection=True,
    export_apply=True,
    export_animations=False,
    export_skins=False,
)
print(f"  Exported: {HEAD_OUTPUT} ({os.path.getsize(HEAD_OUTPUT) // 1024} KB)")

print("\n=== Done ===")