PROCITY/pipeline/merge_ped.py

"""Merge a skinned ped rig's sub-meshes to ONE draw: join meshes → pack all materials into a single
atlas (bake) → one material → one primitive. Skinning (joints/weights) + skeleton are preserved, so
the shared walk.glb/idle.glb clips still retarget. Round-7 E1 (gates D's full-density roster flip).

    BL=/Applications/Blender.app/Contents/MacOS/Blender
    "$BL" --background --python merge_ped.py -- BATCH.json OUT_DIR

BATCH.json = list of {"src": "web/models/peds/x.glb", "name": "x", "atlas": 1024}.
Approach = the bake_lowpoly EMIT trick without the remesh: keep the mesh + armature, give it a
packed atlas UV, route each source material's base-colour→emission via the ORIGINAL UV, bake EMIT
into one atlas image, assign a single material. Falls back to material-merge-only (prims = #mats)
if a ped has no textures to bake.
"""
import bpy, sys, os, json

ARGV = sys.argv[sys.argv.index("--") + 1:]
BATCH, OUT_DIR = ARGV[0], ARGV[1]
os.makedirs(OUT_DIR, exist_ok=True)


def wipe():
    for o in list(bpy.data.objects):
        bpy.data.objects.remove(o, do_unlink=True)
    for coll in (bpy.data.meshes, bpy.data.materials, bpy.data.images, bpy.data.armatures):
        for d in list(coll):
            if d.users == 0:
                coll.remove(d)


def import_join(src):
    bpy.ops.import_scene.gltf(filepath=src)
    meshes = [o for o in bpy.data.objects if o.type == 'MESH']
    bpy.ops.object.select_all(action='DESELECT')
    for m in meshes:
        m.select_set(True)
    bpy.context.view_layer.objects.active = meshes[0]
    if len(meshes) > 1:
        bpy.ops.object.join()                      # keeps vertex groups (weights) + armature modifier
    return bpy.context.view_layer.objects.active


def route_basecolor_to_emission(obj, orig_uv):
    """Each material: sample its base-colour texture via the ORIGINAL uv and feed Emission, so a
    bake on the atlas UV captures albedo. Add an explicit UVMap node so the atlas UV (active during
    bake) doesn't hijack the texture lookup."""
    for m in obj.data.materials:
        if not m or not m.use_nodes:
            continue
        nt = m.node_tree
        b = nt.nodes.get("Principled BSDF")
        if not b:
            continue
        uvn = nt.nodes.new("ShaderNodeUVMap"); uvn.uv_map = orig_uv
        for n in nt.nodes:
            if n.type == 'TEX_IMAGE':
                nt.links.new(uvn.outputs["UV"], n.inputs["Vector"])
        bc = b.inputs.get("Base Color")
        ec = b.inputs.get("Emission Color") or b.inputs.get("Emission")
        es = b.inputs.get("Emission Strength")
        if ec is None:
            continue
        if bc and bc.is_linked:
            nt.links.new(bc.links[0].from_socket, ec)
        elif bc:
            ec.default_value = bc.default_value
        if es is not None:
            es.default_value = 1.0


def merge_one(spec):
    wipe()
    src = spec["src"] if os.path.isabs(spec["src"]) else os.path.join(os.getcwd(), spec["src"])
    obj = import_join(src)
    prim_before = "n/a"
    mats_before = len(obj.data.materials)
    has_tex = any(m and m.use_nodes and any(n.type == 'TEX_IMAGE' for n in m.node_tree.nodes)
                  for m in obj.data.materials)
    size = spec.get("atlas", 1024)

    if mats_before <= 1 or not has_tex:
        # already single-material (or untextured) — nothing to atlas; export as-is (1 primitive)
        out = _export(obj, spec)
        return {"name": spec["name"], "file": os.path.basename(out), "materials": mats_before,
                "atlased": False, "note": "already single-material" if mats_before <= 1 else "untextured"}

    orig_uv = obj.data.uv_layers.active.name
    atlas_uv = obj.data.uv_layers.new(name="atlas").name
    obj.data.uv_layers.active = obj.data.uv_layers[atlas_uv]
    # pack every island of the atlas UV into non-overlapping 0-1 regions (separates the per-material
    # islands so their baked pixels don't collide)
    bpy.ops.object.mode_set(mode='EDIT')
    bpy.ops.mesh.select_all(action='SELECT')
    bpy.ops.uv.select_all(action='SELECT')
    try:
        bpy.ops.uv.pack_islands(rotate=False, margin=0.003)
    except TypeError:
        bpy.ops.uv.pack_islands(margin=0.003)
    bpy.ops.object.mode_set(mode='OBJECT')

    route_basecolor_to_emission(obj, orig_uv)

    img = bpy.data.images.new("ped_atlas", size, size)
    for m in obj.data.materials:
        if m and m.use_nodes:
            tn = m.node_tree.nodes.new("ShaderNodeTexImage"); tn.image = img
            m.node_tree.nodes.active = tn           # bake target per material

    scn = bpy.context.scene
    scn.render.engine = 'CYCLES'
    try:
        scn.cycles.device = 'GPU'
    except Exception:
        pass
    scn.cycles.samples = 4
    scn.render.bake.margin = max(4, size // 128)
    scn.render.bake.use_selected_to_active = False
    obj.data.uv_layers.active = obj.data.uv_layers[atlas_uv]
    bpy.ops.object.select_all(action='DESELECT')
    obj.select_set(True)
    bpy.context.view_layer.objects.active = obj
    bpy.ops.object.bake(type='EMIT', use_clear=True)

    # single atlas material, atlas UV only
    newmat = bpy.data.materials.new("ped_merged"); newmat.use_nodes = True
    nt = newmat.node_tree
    bsdf = nt.nodes.get("Principled BSDF")
    tex = nt.nodes.new("ShaderNodeTexImage"); tex.image = img
    nt.links.new(tex.outputs["Color"], bsdf.inputs["Base Color"])
    obj.data.materials.clear()
    obj.data.materials.append(newmat)
    for uv in list(obj.data.uv_layers):
        if uv.name != atlas_uv:
            obj.data.uv_layers.remove(uv)

    out = _export(obj, spec)
    return {"name": spec["name"], "file": os.path.basename(out), "materials_before": mats_before,
            "atlased": True, "atlas_px": size}


def _export(obj, spec):
    out = os.path.join(OUT_DIR, os.path.basename(spec["src"]))
    bpy.ops.object.select_all(action='DESELECT')
    obj.select_set(True)
    # select the armature too so the skin exports
    for o in bpy.data.objects:
        if o.type == 'ARMATURE':
            o.select_set(True)
    bpy.ops.export_scene.gltf(filepath=out, export_format='GLB', use_selection=True,
                              export_yup=True, export_skins=True, export_animations=False,
                              export_image_format='WEBP', export_image_quality=90)
    return out


def main():
    specs = json.load(open(BATCH))
    results = []
    for spec in specs:
        try:
            r = merge_one(spec)
            results.append(r)
            print(f"OK  {r['file']:32} atlased={r.get('atlased')} mats={r.get('materials_before', r.get('materials'))}")
        except Exception as e:
            print(f"ERR {spec.get('name')}: {e}")
            results.append({"name": spec.get("name"), "error": str(e)})
    json.dump({"results": results}, open(os.path.join(OUT_DIR, "_merge_results.json"), "w"), indent=2)
    print(f"MERGE_DONE ok={sum('error' not in r for r in results)}/{len(specs)}")


main()