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

"""
VALIDATED: Segment a Quaternius body mesh into bone-group regions.

Each segment becomes a separate mesh, skinned to the same skeleton,
exportable independently. Segments map to clothing coverage zones
and can be selectively hidden at runtime.

Segments:
    head        — Head bone
    neck        — neck_01
    torso       — spine_01, spine_02, spine_03, pelvis
    arm_upper_l — clavicle_l, upperarm_l
    arm_upper_r — clavicle_r, upperarm_r
    arm_lower_l — lowerarm_l
    arm_lower_r — lowerarm_r
    hand_l      — hand_l, all finger bones _l
    hand_r      — hand_r, all finger bones _r
    leg_upper_l — thigh_l
    leg_upper_r — thigh_r
    leg_lower_l — calf_l
    leg_lower_r — calf_r
    foot_l      — foot_l, ball_l, ball_leaf_l
    foot_r      — foot_r, ball_r, ball_leaf_r

Run via:
    reach blender run \
        spikes/quaternius-aesthetic/scripts/blender/segment_body.py \
        -- <input.gltf> <output_dir/>

Output:
    <output_dir>/seg_head.glb
    <output_dir>/seg_neck.glb
    <output_dir>/seg_torso.glb
    ... etc
    <output_dir>/armature.glb   (armature only, no body mesh)
"""

import bpy
import sys
import os

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

if len(argv) < 2:
    print("Usage: -- <input.gltf> <output_dir/>")
    sys.exit(1)

INPUT_PATH = argv[0]
OUTPUT_DIR = argv[1]
os.makedirs(OUTPUT_DIR, exist_ok=True)

# Bone-to-segment mapping. Each segment is a list of bone names.
# A vertex belongs to the segment whose bones have the highest combined weight.
SEGMENTS = {
    "head": ["Head"],
    "neck": ["neck_01"],
    "torso": ["spine_01", "spine_02", "spine_03", "pelvis", "root"],
    "arm_upper_l": ["clavicle_l", "upperarm_l"],
    "arm_upper_r": ["clavicle_r", "upperarm_r"],
    "arm_lower_l": ["lowerarm_l"],
    "arm_lower_r": ["lowerarm_r"],
    "hand_l": ["hand_l", "index_01_l", "index_02_l", "index_03_l", "index_04_leaf_l",
                "middle_01_l", "middle_02_l", "middle_03_l", "middle_04_leaf_l",
                "pinky_01_l", "pinky_02_l", "pinky_03_l", "pinky_04_leaf_l",
                "ring_01_l", "ring_02_l", "ring_03_l", "ring_04_leaf_l",
                "thumb_01_l", "thumb_02_l", "thumb_03_l", "thumb_04_leaf_l"],
    "hand_r": ["hand_r", "index_01_r", "index_02_r", "index_03_r", "index_04_leaf_r",
                "middle_01_r", "middle_02_r", "middle_03_r", "middle_04_leaf_r",
                "pinky_01_r", "pinky_02_r", "pinky_03_r", "pinky_04_leaf_r",
                "ring_01_r", "ring_02_r", "ring_03_r", "ring_04_leaf_r",
                "thumb_01_r", "thumb_02_r", "thumb_03_r", "thumb_04_leaf_r"],
    "leg_upper_l": ["thigh_l"],
    "leg_upper_r": ["thigh_r"],
    "leg_lower_l": ["calf_l"],
    "leg_lower_r": ["calf_r"],
    "foot_l": ["foot_l", "ball_l", "ball_leaf_l"],
    "foot_r": ["foot_r", "ball_r", "ball_leaf_r"],
}

print("=== Segment body mesh ===")
print(f"  Input:  {INPUT_PATH}")
print(f"  Output: {OUTPUT_DIR}")
print(f"  Segments: {len(SEGMENTS)}")


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_vertex_segment(mesh_obj, vert_index, vg_name_to_segment):
    """Determine which segment a vertex belongs to based on highest combined bone weight."""
    segment_weights = {}
    for g in mesh_obj.data.vertices[vert_index].groups:
        vg = mesh_obj.vertex_groups[g.group]
        seg = vg_name_to_segment.get(vg.name)
        if seg:
            segment_weights[seg] = segment_weights.get(seg, 0.0) + g.weight

    if not segment_weights:
        return "torso"  # fallback for unweighted verts
    return max(segment_weights, key=segment_weights.get)


def get_neighbor_verts(mesh_obj, vert_indices):
    """Find all vertices connected to the given set by edges (1-ring border)."""
    idx_set = set(vert_indices)
    neighbors = set()
    for edge in mesh_obj.data.edges:
        v0, v1 = edge.vertices[0], edge.vertices[1]
        if v0 in idx_set and v1 not in idx_set:
            neighbors.add(v1)
        elif v1 in idx_set and v0 not in idx_set:
            neighbors.add(v0)
    return neighbors


def extract_segment(mesh_obj, vert_indices, segment_name, armature):
    """
    Create a new mesh from a subset of vertices plus a 1-ring border overlap.
    The overlap ensures adjacent segments share boundary geometry, eliminating
    visible seams at segment boundaries.
    """
    if not vert_indices:
        return None

    # Add 1-ring neighbor vertices as overlap border
    border = get_neighbor_verts(mesh_obj, vert_indices)
    keep_set = set(vert_indices) | border

    # Duplicate the full mesh
    bpy.ops.object.select_all(action='DESELECT')
    mesh_obj.select_set(True)
    mesh_obj.hide_set(False)
    bpy.context.view_layer.objects.active = mesh_obj
    bpy.ops.object.duplicate()
    seg_obj = bpy.context.active_object
    seg_obj.name = "seg_%s" % segment_name

    # Select only the vertices NOT in this segment (+ border) and delete them
    bpy.ops.object.mode_set(mode='EDIT')
    bpy.ops.mesh.select_all(action='DESELECT')
    bpy.ops.object.mode_set(mode='OBJECT')

    for v in seg_obj.data.vertices:
        v.select = v.index not in keep_set

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

    remaining = len(seg_obj.data.vertices)
    if remaining == 0:
        bpy.data.objects.remove(seg_obj)
        return None

    return seg_obj


# --- Main ---

clear_scene()

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

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

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

body_mesh = max(meshes, key=lambda m: len(m.data.vertices))
small_meshes = [m for m in meshes if m != body_mesh]

print(f"  Armature: {armature.name} ({len(armature.data.bones)} bones)")
print(f"  Body mesh: {body_mesh.name} ({len(body_mesh.data.vertices)} verts)")
for m in small_meshes:
    print(f"  Extra mesh: {m.name} ({len(m.data.vertices)} verts)")

# Build reverse map: bone name -> segment name
vg_name_to_segment = {}
for seg_name, bones in SEGMENTS.items():
    for bone in bones:
        vg_name_to_segment[bone] = seg_name

# Classify every vertex
print("\nStep 2: Classifying vertices...")
vert_segments = {}  # segment_name -> [vert_indices]
for v in body_mesh.data.vertices:
    seg = get_vertex_segment(body_mesh, v.index, vg_name_to_segment)
    if seg not in vert_segments:
        vert_segments[seg] = []
    vert_segments[seg].append(v.index)

for seg_name in sorted(vert_segments.keys()):
    print(f"  {seg_name}: {len(vert_segments[seg_name])} verts")

# Extract each segment
print("\nStep 3: Extracting segments...")
segment_objects = {}

for seg_name in SEGMENTS:
    indices = vert_segments.get(seg_name, [])
    if not indices:
        print(f"  {seg_name}: SKIP (no vertices)")
        continue

    seg_obj = extract_segment(body_mesh, indices, seg_name, armature)
    if seg_obj:
        segment_objects[seg_name] = seg_obj
        print(f"  {seg_name}: {len(seg_obj.data.vertices)} verts")
    else:
        print(f"  {seg_name}: FAILED")

# Hide the original body mesh
body_mesh.hide_set(True)

# Export each segment with the armature
print("\nStep 4: Exporting segments...")
for seg_name, seg_obj in segment_objects.items():
    # Hide all segments except this one
    for other_name, other_obj in segment_objects.items():
        other_obj.hide_set(other_name != seg_name)
    for m in small_meshes:
        m.hide_set(True)

    bpy.ops.object.select_all(action='DESELECT')
    armature.select_set(True)
    seg_obj.select_set(True)
    bpy.context.view_layer.objects.active = armature

    output_path = os.path.join(OUTPUT_DIR, "seg_%s.glb" % seg_name)
    bpy.ops.export_scene.gltf(
        filepath=output_path,
        export_format='GLB',
        use_selection=True,
        export_apply=False,
        export_animations=False,
        export_skins=True,
    )
    size_kb = os.path.getsize(output_path) // 1024
    print(f"  seg_{seg_name}.glb ({size_kb} KB)")

# Export eyes/eyebrows as separate segments too
print("\nStep 5: Exporting extra meshes...")
for m in small_meshes:
    for other_obj in segment_objects.values():
        other_obj.hide_set(True)
    for other_m in small_meshes:
        other_m.hide_set(other_m != m)

    bpy.ops.object.select_all(action='DESELECT')
    armature.select_set(True)
    m.select_set(True)
    m.hide_set(False)
    bpy.context.view_layer.objects.active = armature

    clean_name = m.name.lower().replace(" ", "_")
    output_path = os.path.join(OUTPUT_DIR, "seg_%s.glb" % clean_name)
    bpy.ops.export_scene.gltf(
        filepath=output_path,
        export_format='GLB',
        use_selection=True,
        export_apply=False,
        export_animations=False,
        export_skins=True,
    )
    size_kb = os.path.getsize(output_path) // 1024
    print(f"  seg_{clean_name}.glb ({size_kb} KB)")

# Export the armature alone (for animation loading)
print("\nStep 6: Exporting armature...")
for obj in segment_objects.values():
    obj.hide_set(True)
for m in small_meshes:
    m.hide_set(True)
body_mesh.hide_set(True)

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

print("\n=== Done ===")
for seg_name in SEGMENTS:
    path = os.path.join(OUTPUT_DIR, "seg_%s.glb" % seg_name)
    status = "[OK]" if os.path.exists(path) else "[MISSING]"
    print(f"  {status} seg_{seg_name}.glb")