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 nonmanifold_ratio(obj):
    """Fraction of vertices on a non-manifold edge — the number that predicts whether COLLAPSE
    decimation can work at all. Measured, not guessed: see remesh_shellsoup()."""
    n = len(obj.data.vertices)
    if not n:
        return 0.0
    bpy.ops.object.mode_set(mode='EDIT')
    bpy.ops.mesh.select_all(action='DESELECT')
    bpy.ops.mesh.select_non_manifold()
    bpy.ops.object.mode_set(mode='OBJECT')
    return sum(1 for v in obj.data.vertices if v.select) / n


def remesh_shellsoup(obj, budget, voxel=None, thick=0.0):
    """Voxel-remesh a shell-soup mesh to a manifold surface, decimate to budget, and carry the
    source's vertex colours across.

    [R41 §41.2] `art_incoming`'s 48 pub props are trimesh output at ~200k tris that COLLAPSE
    decimation cannot touch: loungeChair is 94% non-manifold (376,076 of 399,542 verts) and a
    ratio-0.01 decimate moved it 1,319,836 -> 1,274,119, i.e. 3%. topiary behaved the same at 93%
    non-manifold; pokie, at 29%, decimated 198k -> 47k normally. Decimate cannot merge across
    overlapping shells, so the fix is to stop trying: rebuild the topology with a voxel remesh
    (OpenVDB, always manifold), then decimate THAT, then transfer the colour back.

    This is the same "throw the topology away" move as bake_lowpoly.py, minus the Cycles bake —
    these props carry COLOR_0 and no textures, so a vertex-colour transfer keeps them at one
    material and zero texture bytes instead of adding a 1024px atlas each.
    """
    src = obj
    bpy.ops.object.select_all(action='DESELECT')
    src.select_set(True)
    bpy.context.view_layer.objects.active = src
    bpy.ops.object.duplicate()
    lp = bpy.context.view_layer.objects.active          # the copy becomes the low-poly
    lp.name = "LP"

    dims = lp.dimensions
    v = voxel or max(max(dims) / 150.0, 1e-4)

    # A voxel remesh is a signed-distance field: it needs a VOLUME. Feed it a thin open shell and
    # it returns a lattice of holes — the keg (a thin-walled scan) came back as a wire cage at
    # every voxel size tried, while solid props like the couch remeshed cleanly. Giving the shell
    # a wall thicker than one voxel first is what makes the field well-defined.
    if thick:
        sm = lp.modifiers.new("solid", 'SOLIDIFY')
        sm.thickness = thick if thick > 0 else v * 2.5
        sm.offset = 0.0
        sm.use_even_offset = False
        bpy.ops.object.modifier_apply(modifier=sm.name)

    rm = lp.modifiers.new("remesh", 'REMESH')
    rm.mode = 'VOXEL'
    rm.voxel_size = v
    rm.adaptivity = 0.0
    bpy.ops.object.modifier_apply(modifier=rm.name)

    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 > budget:
        dm = lp.modifiers.new("dec", 'DECIMATE')
        dm.ratio = max(0.002, budget / n)
        bpy.ops.object.modifier_apply(modifier=dm.name)

    # Carry colour (and UVs, if the source is textured) across from the dense original by
    # nearest-face interpolation. The destination layers must be CREATED first with
    # `datalayout_transfer` — a Data Transfer modifier applied without it silently transfers
    # nothing and the prop exports colourless, which is exactly how the first run of this came
    # back "colour LOST" on all six.
    transferred = False
    src_cols = list(src.data.color_attributes)
    src_uvs = list(src.data.uv_layers)
    if src_cols or src_uvs:
        try:
            dt = lp.modifiers.new("xfer", 'DATA_TRANSFER')
            dt.object = src
            loops, verts = set(), set()
            if src_cols:
                dom = (src.data.color_attributes.active_color or src_cols[0]).domain
                (loops if dom == 'CORNER' else verts).add(
                    'COLOR_CORNER' if dom == 'CORNER' else 'COLOR_VERTEX')
            if src_uvs:
                loops.add('UV')
            if loops:
                dt.use_loop_data = True
                dt.data_types_loops = loops
                dt.loop_mapping = 'POLYINTERP_NEAREST'
            if verts:
                dt.use_vert_data = True
                dt.data_types_verts = verts
                dt.vert_mapping = 'POLYINTERP_NEAREST'
            bpy.ops.object.datalayout_transfer(modifier=dt.name)     # create dest layers
            bpy.ops.object.modifier_apply(modifier=dt.name)
            transferred = bool(lp.data.color_attributes) if src_cols else bool(lp.data.uv_layers)
        except Exception as e:
            print(f"    colour/UV transfer failed: {e}")

    # the voxel remesh can leave degenerate faces; Blender warns "mesh is not valid" on export
    lp.data.validate(verbose=False)

    lp_mesh_name = lp.data.name
    bpy.data.objects.remove(src, do_unlink=True)
    bpy.ops.object.select_all(action='DESELECT')
    lp.select_set(True)
    bpy.context.view_layer.objects.active = lp
    print(f"    remesh voxel={v:.4f} -> {n} tris -> {len(lp.data.polygons)} "
          f"(colour {'transferred' if transferred else 'LOST'}) [{lp_mesh_name}]")
    return lp, transferred


def wire_vertex_colour(obj):
    """Route the active colour attribute into Base Color on every material.

    The glTF exporter only writes COLOR_0 when a material actually reads the attribute
    (`export_vertex_color='MATERIAL'`), and these sources arrive with NO material block at all —
    Blender invents a plain grey one on import. Without this the remeshed props would export
    colourless and every thumbnail would be grey.
    """
    if not obj.data.color_attributes:
        return False
    cname = obj.data.color_attributes.active_color.name \
        if obj.data.color_attributes.active_color else obj.data.color_attributes[0].name
    if not obj.data.materials:
        m = bpy.data.materials.new("vcol")
        m.use_nodes = True
        obj.data.materials.append(m)
    for m in obj.data.materials:
        if not m or not m.use_nodes:
            continue
        nt = m.node_tree
        bsdf = nt.nodes.get("Principled BSDF")
        if not bsdf:
            continue
        n = nt.nodes.new("ShaderNodeVertexColor")
        n.layer_name = cname
        nt.links.new(n.outputs["Color"], bsdf.inputs["Base Color"])
        if "Specular IOR Level" in bsdf.inputs:
            bsdf.inputs["Specular IOR Level"].default_value = 0.2
        bsdf.inputs["Roughness"].default_value = 0.72
    return True


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)

    # shell-soup path: rebuild topology before the decimate loop, which cannot touch it
    nm = 0.0
    remeshed = vcol = False
    if spec.get("remesh"):
        nm = nonmanifold_ratio(obj)
        obj, vcol = remesh_shellsoup(obj, spec.get("tris", 5000), spec.get("voxel"),
                                     spec.get("shell_thickness", 0.0))
        remeshed = True        # 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))
    if rx or ry or rz or spec.get("yaw"):
        # Select+activate explicitly: transform_apply acts on the SELECTION, and after clean_mesh
        # (or a remesh) the object we mean is not guaranteed to still be it. A silently-skipped
        # yaw is invisible in the log and ships a prop facing the wrong way.
        bpy.ops.object.select_all(action='DESELECT')
        obj.select_set(True)
        bpy.context.view_layer.objects.active = obj
        obj.rotation_mode = 'XYZ'
        before = tuple(round(v, 4) for v in obj.dimensions)
        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)
        after = tuple(round(v, 4) for v in obj.dimensions)
        if before == after and (spec.get("yaw", 0) % 180) not in (0,):
            print(f"    WARNING: yaw/up requested but dimensions unchanged {before} — not applied")
        else:
            print(f"    oriented {before} -> {after}")

    # 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)
    if spec.get("remesh") or spec.get("vertex_colour"):
        vcol = wire_vertex_colour(obj) or vcol

    name = f"procity_{spec['category']}_{spec['id']}_01.glb"
    out = os.path.join(OUT_DIR, name)
    exp = dict(filepath=out, export_format='GLB', export_yup=True, export_apply=True,
               export_image_format='WEBP', export_image_quality=85)
    try:                                  # Blender >=4.2 gates COLOR_0 behind this enum
        bpy.ops.export_scene.gltf(**exp, export_vertex_color='ACTIVE')
    except TypeError:
        bpy.ops.export_scene.gltf(**exp)
    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,
        "remeshed": remeshed, "nonmanifold": round(nm, 3), "vertex_colour": vcol,
        "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()