PROCITY/pipeline/bake_lowpoly.py

"""PROCITY hero-prop bake-to-lowpoly — turn a heavy TRELLIS shell-soup GLB into a game-ready
≤8k-tri prop with its base colour BAKED onto fresh UVs.

Runs headless on this box's Blender:
    BL=/Applications/Blender.app/Contents/MacOS/Blender
    "$BL" --background --python bake_lowpoly.py -- BATCH.json OUT_DIR [THUMB_DIR]

Why bake (R3's lesson): raw TRELLIS meshes are dense multi-shell scans; plain decimate /
meshopt-simplify floors at 25–77k because they can't merge across the shells. So we throw the
topology away and rebuild it:
  1. import the normalized HP (already metres / +Y up / base-origin / textured)  → bake SOURCE
  2. voxel-remesh a copy  → one watertight manifold surface (no more shell-soup)
  3. decimate the remesh to the tri budget (≤8k)  → LP
  4. smart-UV unwrap LP
  5. Cycles bake DIFFUSE(COLOR) HP → LP's new UVs  → a single baked base-colour texture
  6. export LP GLB (WebP, embedded, self-illuminated a touch like normalize.py) + 256 thumbnail

BATCH.json: [ {"src":"pipeline/_normalized/procity_fit_drinks_fridge_01.glb",
               "id":"drinks_fridge", "category":"fit",
               "voxel":0.012, "tris":8000, "img":1024} , ... ]
"""
import bpy, sys, os, json, math
from mathutils import Vector

ARGV = sys.argv[sys.argv.index("--") + 1:]
BATCH, OUT_DIR = ARGV[0], (ARGV[1] if len(ARGV) > 1 else "/tmp/procity_bake")
THUMB_DIR = ARGV[2] if len(ARGV) > 2 else os.path.join(OUT_DIR, "thumbs")
os.makedirs(OUT_DIR, exist_ok=True); os.makedirs(THUMB_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):
        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()
    obj = bpy.context.view_layer.objects.active
    bpy.ops.object.transform_apply(location=True, rotation=True, scale=True)
    return obj


def bounds(obj):
    bb = [obj.matrix_world @ Vector(c) for c in obj.bound_box]
    mn = Vector((min(v.x for v in bb), min(v.y for v in bb), min(v.z for v in bb)))
    mx = Vector((max(v.x for v in bb), max(v.y for v in bb), max(v.z for v in bb)))
    return mn, mx


def make_lowpoly(hp, voxel, tris):
    """Duplicate HP → voxel remesh → decimate to budget → smart UV. Returns the LP object."""
    bpy.ops.object.select_all(action='DESELECT')
    hp.select_set(True); bpy.context.view_layer.objects.active = hp
    bpy.ops.object.duplicate()
    lp = bpy.context.view_layer.objects.active
    lp.name = "LP"
    # strip HP's materials off the LP — it gets one fresh baked material
    lp.data.materials.clear()
    rm = lp.modifiers.new("remesh", 'REMESH')
    rm.mode = 'VOXEL'; rm.voxel_size = voxel; rm.adaptivity = 0.0
    bpy.ops.object.modifier_apply(modifier=rm.name)
    # triangulate so len(polygons) == triangle count → the decimate ratio targets triangles honestly
    bpy.ops.object.mode_set(mode='EDIT')
    bpy.ops.mesh.select_all(action='SELECT')
    bpy.ops.mesh.quads_convert_to_tris()
    bpy.ops.object.mode_set(mode='OBJECT')
    n = len(lp.data.polygons)
    if n > tris:
        dm = lp.modifiers.new("dec", 'DECIMATE')
        dm.ratio = max(0.005, tris / n)
        bpy.ops.object.modifier_apply(modifier=dm.name)
    # UVs
    bpy.ops.object.mode_set(mode='EDIT')
    bpy.ops.mesh.select_all(action='SELECT')
    bpy.ops.uv.smart_project(island_margin=0.02, angle_limit=math.radians(66))
    bpy.ops.object.mode_set(mode='OBJECT')
    return lp


def _base_color_to_emission(hp):
    """Route every HP material's base-colour (texture or flat value) into its Emission and mute the
    BSDF, so a bake_type='EMIT' pass reads PURE ALBEDO regardless of the source material setup
    (TRELLIS and fal-Hunyuan wire their materials differently; DIFFUSE(COLOR) mis-reads fal's)."""
    for m in hp.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
        bc, ec = b.inputs.get("Base Color"), (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 bake_color(hp, lp, img_size):
    _base_color_to_emission(hp)
    img = bpy.data.images.new("baked", img_size, img_size)
    mat = bpy.data.materials.new("baked_mat"); mat.use_nodes = True
    nt = mat.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"])
    lp.data.materials.clear(); lp.data.materials.append(mat)
    nt.nodes.active = tex                      # the active image node = bake target

    scn = bpy.context.scene
    scn.render.engine = 'CYCLES'
    try:
        scn.cycles.device = 'GPU'
    except Exception:
        pass
    scn.cycles.samples = 4
    scn.render.bake.use_selected_to_active = True
    scn.render.bake.cage_extrusion = 0.06
    scn.render.bake.max_ray_distance = 0.12
    scn.render.bake.margin = max(4, img_size // 128)

    bpy.ops.object.select_all(action='DESELECT')
    hp.select_set(True); lp.select_set(True)
    bpy.context.view_layer.objects.active = lp   # active = bake TARGET
    bpy.ops.object.bake(type='EMIT', use_clear=True)   # HP base-colour routed to emission = albedo

    # self-illuminate a touch so it reads in any viewer (matches normalize.py bright)
    ec = bsdf.inputs.get("Emission Color") or bsdf.inputs.get("Emission")
    if ec is not None:
        nt.links.new(tex.outputs["Color"], ec)
        es = bsdf.inputs.get("Emission Strength")
        if es is not None:
            es.default_value = 0.28
    return img


def base_origin(obj):
    mn, mx = bounds(obj)
    obj.location -= Vector(((mn.x + mx.x) / 2, (mn.y + mx.y) / 2, mn.z))
    bpy.ops.object.transform_apply(location=True)


def render_thumb(lp, png_name):
    scn = bpy.context.scene
    mn, mx = bounds(lp); ctr = (mn + mx) / 2
    size = max((mx - mn).x, (mx - mn).y, (mx - mn).z) or 1.0
    for o in list(bpy.data.objects):
        if o.type in ('CAMERA', 'LIGHT'):
            bpy.data.objects.remove(o, do_unlink=True)
    w = bpy.data.worlds.get("tw") or bpy.data.worlds.new("tw")
    scn.world = w; w.use_nodes = True
    w.node_tree.nodes["Background"].inputs[1].default_value = 1.05
    bpy.ops.object.light_add(type='AREA', location=(ctr.x + size, ctr.y - size, mx.z + size))
    bpy.context.object.data.energy = 800 * size; bpy.context.object.data.size = 6
    bpy.ops.object.camera_add(); cam = bpy.context.object; scn.camera = cam
    r = size * 1.9
    cam.location = Vector((ctr.x + r * 0.8, ctr.y - r, ctr.z + size * 0.6))
    cam.rotation_euler = (ctr - cam.location).to_track_quat('-Z', 'Y').to_euler()
    scn.render.engine = 'BLENDER_WORKBENCH'
    scn.display.shading.light = 'STUDIO'; scn.display.shading.color_type = 'TEXTURE'
    scn.display.shading.show_cavity = True
    scn.view_settings.view_transform = 'Standard'
    scn.render.film_transparent = True
    scn.render.resolution_x = scn.render.resolution_y = 256
    p = os.path.join(THUMB_DIR, png_name); scn.render.filepath = p
    bpy.ops.render.render(write_still=True)
    return p


def bake_one(spec):
    wipe()
    src = spec["src"] if os.path.isabs(spec["src"]) else os.path.join(os.getcwd(), spec["src"])
    hp = import_join(src)
    tri_before = len(hp.data.polygons)
    lp = make_lowpoly(hp, spec.get("voxel", 0.02), spec.get("tris", 8000))
    bake_color(hp, lp, spec.get("img", 1024))
    bpy.data.objects.remove(hp, do_unlink=True)          # HP done; export LP only
    base_origin(lp)
    lp.select_set(True); bpy.context.view_layer.objects.active = lp
    name = f"procity_{spec['category']}_{spec['id']}_01.glb"
    out = os.path.join(OUT_DIR, name)
    bpy.ops.export_scene.gltf(filepath=out, export_format='GLB', export_yup=True,
                              use_selection=True, export_image_format='WEBP', export_image_quality=88)
    thumb = render_thumb(lp, name.replace(".glb", ".png"))
    mn, mx = bounds(lp)
    return {"file": name, "tri_before": tri_before, "tri_after": len(lp.data.polygons),
            "size_m": [round(mx.x - mn.x, 3), round(mx.z - mn.z, 3), round(mx.y - mn.y, 3)],
            "footprint": [round(mx.x - mn.x, 2), round(mx.y - mn.y, 2)], "thumb": thumb, "out": out}


def main():
    specs = json.load(open(BATCH))
    results, errors = [], []
    for spec in specs:
        try:
            r = bake_one(spec)
            results.append(r)
            print(f"OK  {r['file']:36} {r['tri_before']:>6}→{r['tri_after']:>5} tris  "
                  f"footprint {r['footprint']} h {r['size_m'][1]}")
        except Exception as e:
            import traceback; traceback.print_exc()
            errors.append({"id": spec.get("id"), "error": str(e)})
            print(f"ERR {spec.get('id')}: {e}")
    json.dump({"results": results, "errors": errors},
              open(os.path.join(OUT_DIR, "_bake_results.json"), "w"), indent=2)
    print(f"\nBAKE_DONE ok={len(results)} err={len(errors)} → {OUT_DIR}")


main()