PROCITY/pipeline/normalize.py

"""PROCITY GLB normalizer — enforce the house GLB law on library assets.

Runs headless on ultra (Blender at /Applications/Blender.app/Contents/MacOS/Blender):

    BL=/Applications/Blender.app/Contents/MacOS/Blender
    "$BL" --background --python normalize.py -- BATCH.json OUT_DIR [THUMB_DIR]

Processes every entry in BATCH.json in ONE Blender session (clean scene between
assets — no per-asset relaunch). For each asset it:

  * imports (Blender applies the glTF Y-up convention → correct world orientation)
  * joins to a single mesh, applies transforms
  * optional up-axis correction + yaw so the asset stands up and faces -Z
  * base-origin: centre in X/Y, sit on the floor (min Z = 0)          [house law]
  * height-normalises to a real-world metre height                    [house law]
  * decimates to a triangle budget (props <=5k)                       [house law]
  * optional self-illuminate (TRELLIS/dark-bake assets)   (finish_glb.py:88)
  * downsizes every texture to <=1024px and exports them as WebP      [house law]
  * exports GLB, +Y up, embedded textures, NO Draco                   [house law]
  * renders a 256px hero thumbnail (workbench, textured)              [step 2]

Output names: procity_<category>_<id>_01.glb  (e.g. procity_fit_record_crate_01.glb)

The heavy finish logic (solidify/decimate/bright/base-origin) is lifted verbatim
from MESHGOD/scripts/finish_glb.py — this is the same law, extended for the
PROCITY batch + texture-webp + thumbnail, per the lane rule "extend, don't fork".

BATCH.json schema (list of):
  { "src": "shop-fittings/balcao.glb",   # relative to --root (default ~/Documents/3D=models)
    "id": "counter", "category": "fit",  # → procity_fit_counter_01.glb
    "height_m": 1.1,                      # target height (0 = keep source scale)
    "tris": 3000,                         # triangle budget
    "up": "none",                         # none|x+90|x-90|y+90|y-90  (stand a lying asset up)
    "yaw": 0,                             # degrees about up-axis, so front faces -Z
    "solidify": false,                    # cap open bottoms (TRELLIS scans only)
    "bright": false }                     # self-illuminate dark-baked albedo
"""
import bpy, sys, os, math, json
from mathutils import Vector

ARGV = sys.argv[sys.argv.index("--") + 1:]
BATCH = ARGV[0]
OUT_DIR = ARGV[1] if len(ARGV) > 1 else "/tmp/procity_norm"
THUMB_DIR = ARGV[2] if len(ARGV) > 2 else os.path.join(OUT_DIR, "thumbs")
ROOT = os.environ.get("PROCITY_MODELS_ROOT", os.path.expanduser("~/Documents/3D=models"))
os.makedirs(OUT_DIR, exist_ok=True)
os.makedirs(THUMB_DIR, exist_ok=True)

UP_ROT = {                     # extra rotation to stand a mis-authored asset upright
    "none": (0, 0, 0), "x+90": (math.radians(90), 0, 0), "x-90": (math.radians(-90), 0, 0),
    "y+90": (0, math.radians(90), 0), "y-90": (0, math.radians(-90), 0),
    "z+90": (0, 0, math.radians(90)), "z-90": (0, 0, math.radians(-90)),
}


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


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 import_join(src):
    ext = os.path.splitext(src)[1].lower()
    if ext in (".glb", ".gltf"):
        bpy.ops.import_scene.gltf(filepath=src)
    elif ext == ".fbx":
        bpy.ops.import_scene.fbx(filepath=src)
    elif ext == ".obj":
        bpy.ops.wm.obj_import(filepath=src)
    else:
        raise ValueError("unsupported: " + ext)
    meshes = [o for o in bpy.data.objects if o.type == 'MESH']
    if not meshes:
        raise ValueError("no meshes")
    # BAKE ancestor transforms into each mesh, then drop non-mesh nodes. glТF imports often park
    # the mesh under a parent empty carrying unit-conversion scale / pivot offset — joining +
    # transform_apply alone leaves that parent node in the export (asset ships off-origin & scaled,
    # violating "origin at base"). parent_clear KEEP_TRANSFORM folds matrix_world into each basis.
    bpy.ops.object.select_all(action='DESELECT')
    for m in meshes:
        m.select_set(True)
    bpy.context.view_layer.objects.active = meshes[0]
    bpy.ops.object.parent_clear(type='CLEAR_KEEP_TRANSFORM')
    for o in list(bpy.data.objects):
        if o.type != 'MESH':                 # empties, armatures, cameras, lights
            bpy.data.objects.remove(o, do_unlink=True)
    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
    obj.modifiers.clear()                    # drop armature/etc modifiers (static props)
    bpy.ops.object.transform_apply(location=True, rotation=True, scale=True)
    return obj


def clean_mesh(obj):
    """Drop stray/loose geometry that inflates the bbox (dumped assets are full of it):
    weld doubles, delete loose verts/edges, then clip statistical outlier verts
    (|coord-median| > 6*MAD on any axis) — the stray points that make a table read 55m."""
    bpy.ops.object.mode_set(mode='EDIT')
    bpy.ops.mesh.select_all(action='SELECT')
    bpy.ops.mesh.remove_doubles(threshold=0.0001)
    bpy.ops.mesh.delete_loose(use_verts=True, use_edges=True, use_faces=False)
    bpy.ops.mesh.quads_convert_to_tris()     # so len(polygons)==triangle count: the ≤5k budget
    bpy.ops.object.mode_set(mode='OBJECT')   # gate + reported tri counts stay honest (glTF triangulates on export anyway)
    vs = obj.data.vertices
    n = len(vs)
    if n < 20:
        return 0
    outliers = []
    for axis in range(3):
        vals = sorted(v.co[axis] for v in vs)
        med = vals[n // 2]
        mad = sorted(abs(x - med) for x in vals)[n // 2] or 1e-6
        lo, hi = med - 6 * mad * 1.4826, med + 6 * mad * 1.4826
        for v in vs:
            if v.co[axis] < lo or v.co[axis] > hi:
                outliers.append(v.index)
    outliers = set(outliers)
    if outliers and len(outliers) < n * 0.03:      # only if they are a tiny minority
        for v in vs:
            v.select = v.index in outliers
        bpy.ops.object.mode_set(mode='EDIT')
        bpy.ops.mesh.delete(type='VERT')
        bpy.ops.object.mode_set(mode='OBJECT')
        return len(outliers)
    return 0


def resize_textures(max_px=1024):
    for img in bpy.data.images:
        if img.size[0] > max_px or img.size[1] > max_px:
            w, h = img.size
            s = max_px / max(w, h)
            img.scale(max(1, int(w * s)), max(1, int(h * s)))


def normalize_one(spec):
    wipe()
    src = os.path.join(ROOT, spec["src"]) if not os.path.isabs(spec["src"]) else spec["src"]
    obj = import_join(src)
    clipped = clean_mesh(obj)                 # welds, drops loose geo, TRIANGULATES
    tri_before = len(obj.data.polygons)        # true source triangle count (post-triangulate)

    # stand upright (if mis-authored), then face -Z
    rx, ry, rz = UP_ROT.get(spec.get("up", "none"), (0, 0, 0))
    obj.rotation_euler = (rx, ry, rz)
    if spec.get("yaw"):
        obj.rotation_euler[2] += math.radians(spec["yaw"])
    bpy.ops.object.transform_apply(rotation=True)

    # solidify: cap open bottoms → watertight, white sole (finish_glb.py:54)
    if spec.get("solidify"):
        bpy.ops.object.mode_set(mode='EDIT')
        bpy.ops.mesh.select_all(action='SELECT')
        bpy.ops.mesh.remove_doubles(threshold=0.0002)
        n0 = len(obj.data.polygons)
        bpy.ops.mesh.fill_holes(sides=0)
        bpy.ops.mesh.normals_make_consistent(inside=False)
        bpy.ops.object.mode_set(mode='OBJECT')
        n1 = len(obj.data.polygons)
        if n1 > n0:
            white = bpy.data.materials.new("sole_white"); white.use_nodes = True
            b = white.node_tree.nodes.get("Principled BSDF")
            if b:
                b.inputs["Base Color"].default_value = (1, 1, 1, 1)
            obj.data.materials.append(white)
            widx = len(obj.data.materials) - 1
            for i in range(n0, n1):
                obj.data.polygons[i].material_index = widx

    # decimate to budget (finish_glb.py:80). TRELLIS meshes are dense multi-shell scans where a
    # single COLLAPSE pass stalls far above target (arcade: 196k→25k at ratio 0.03); iterate —
    # each pass re-targets the budget from the now-smaller count — until within 10% or a pass
    # stops making progress (collapse limit). Library assets already under budget skip the loop.
    budget = spec.get("tris", 5000)
    tris = len(obj.data.polygons)
    passes = 0
    while tris > budget * 1.1 and passes < 5:
        m = obj.modifiers.new('decimate', 'DECIMATE')
        m.ratio = max(0.02, budget / tris)
        bpy.ops.object.modifier_apply(modifier=m.name)
        new = len(obj.data.polygons)
        passes += 1
        if new >= tris * 0.97:        # barely moved — collapse limit reached, don't spin
            tris = new
            break
        tris = new
    tri_after = len(obj.data.polygons)

    # self-illuminate dark bakes (finish_glb.py:88)
    if spec.get("bright"):
        for m in bpy.data.materials:
            if not m.use_nodes:
                continue
            b = m.node_tree.nodes.get("Principled BSDF")
            if not b:
                continue
            bc = b.inputs.get("Base Color"); ec = b.inputs.get("Emission Color")
            es = b.inputs.get("Emission Strength")
            if ec is None:
                continue
            if bc and bc.is_linked:
                m.node_tree.links.new(bc.links[0].from_socket, ec)
            elif bc:
                ec.default_value = bc.default_value
            if es is not None:
                es.default_value = 0.32

    # base-origin + height-normalise (finish_glb.py:109)
    mn, mx = bounds(obj)
    h = spec.get("height_m", 0)
    if h and h > 0:
        cur = (mx.z - mn.z) or 1.0
        obj.scale *= (h / cur)
        bpy.ops.object.transform_apply(scale=True)
        mn, mx = bounds(obj)
    base = Vector(((mn.x + mx.x) / 2, (mn.y + mx.y) / 2, mn.z))
    obj.location -= base
    bpy.ops.object.transform_apply(location=True)
    mn, mx = bounds(obj)

    resize_textures(1024)

    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, export_apply=True,
        export_image_format='WEBP', export_image_quality=85,
    )
    thumb = render_thumb(name.replace(".glb", ".png"))
    return {
        "src": spec["src"], "file": name, "out": out, "thumb": thumb,
        "tri_before": tri_before, "tri_after": tri_after,
        "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)],
    }


def render_thumb(png_name):
    """One 3/4 hero shot, 256px, textured workbench (render_glb.py lighting)."""
    scn = bpy.context.scene
    obj = bpy.context.view_layer.objects.active
    mn, mx = bounds(obj)
    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))
    d = ctr - cam.location
    cam.rotation_euler = d.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 main():
    specs = json.load(open(BATCH))
    results, errors = [], []
    for spec in specs:
        try:
            r = normalize_one(spec)
            results.append(r)
            print(f"OK  {r['file']:34} tris {r['tri_before']}→{r['tri_after']:>5} "
                  f"size {r['size_m']}  footprint {r['footprint']}")
        except Exception as e:
            errors.append({"src": spec.get("src"), "id": spec.get("id"), "error": str(e)})
            print(f"ERR {spec.get('src')}: {e}")
    json.dump({"results": results, "errors": errors},
              open(os.path.join(OUT_DIR, "_results.json"), "w"), indent=2)
    print(f"\nNORMALIZE_DONE ok={len(results)} err={len(errors)} → {OUT_DIR}/_results.json")


main()