PROCITY/pipeline/tint_atlas.py

#!/usr/bin/env python3
"""PROCITY atlas tint — repaint anatomical REGIONS of a baked GLB atlas, in 3D, not in 2D.

Why this exists (R38): `trellis2_mlx`'s PBR bake crushes high-contrast markings. The magpie came
back an all-black bird — a crow — with its white nape, wing bar and tail base gone (AUDIT R38 §4d
proved the loss is the operator's bake: the RMBG cutout and the raw 500k mesh were both checked).
The atlas is a smart-UV bake: ~100 scattered islands, so "hand-tint the white bits" is not a 2D
paint job — you cannot tell which island is a nape.

So the tint is defined in the MESH's own space and projected through the UVs:

  1. rasterize every UV triangle into the atlas grid  ->  texel -> (position, normal) maps
  2. evaluate soft anatomical region masks on those POSITIONS (metres, GLB frame: +Y up)
  3. paint the region colour with the bake's own luminance kept as shading
  4. dilate the result a few texels into the bake's UV padding so no black fringe survives
     bilinear/mip sampling at the island edges
  5. re-embed as WebP and rewrite the GLB — GEOMETRY IS UNTOUCHED, byte for byte. Only
     bufferView 'image' changes, so tris / materials / images / extensions all stay put.

    PY=~/Documents/MODELBEAST/venvs/mflux/bin/python     # numpy + PIL live here
    $PY pipeline/tint_atlas.py magpie --debug   # flat region colours, for looking at
    $PY pipeline/tint_atlas.py magpie           # the real tint
    $PY pipeline/tint_atlas.py magpie --dump-atlas /tmp/x.png   # write the atlas alongside

Region masks are smooth (smoothstep ramps on windows of z / y / |x| / normal), so the paint has a
soft edge in 3D and therefore across every island at once — which is exactly what a 2D blur could
never do here (it would bleed island-to-island across the black background).
"""
import json, os, struct, sys, io
import numpy as np
from PIL import Image

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))


# ---------------------------------------------------------------- soft window helpers
def smoothstep(e0, e1, x):
    t = np.clip((x - e0) / (e1 - e0 + 1e-12), 0.0, 1.0)
    return t * t * (3 - 2 * t)


def window(v, lo_out, lo_in, hi_in, hi_out):
    """1 inside [lo_in, hi_in], ramping to 0 at lo_out / hi_out."""
    return smoothstep(lo_out, lo_in, v) * (1.0 - smoothstep(hi_in, hi_out, v))


# ---------------------------------------------------------------- the region specs
# The AUSTRALIAN MAGPIE (Gymnorhina tibicen), adult, black-backed form — the widespread depiction
# and the one R38 asked for: WHITE NAPE (hindneck, behind a black crown), WHITE SHOULDER/WING BAR
# along the top of the folded wing, WHITE RUMP running into a WHITE TAIL BASE with a black
# terminal band. Beak pale blue-grey with a black tip (the bake already has that and it survives).
#
# Mesh frame, measured off the bytes: x = ±0.0777 lateral · y = 0 -> 0.3583 up ·
# z = -0.2021 (tail tip, drooped to the ground) -> +0.2021 (beak tip). Head/crown z 0.10-0.135,
# neck z 0.055-0.105, body z -0.06..+0.09, legs y<0.09, tail z<-0.06 sloping down to the ground.
MAGPIE = {
    "file": "procity_street_magpie_01.glb",
    "regions": [
        # --- the nape: the broad white collar over the hindneck, stopping short of the black crown
        {"name": "nape", "debug": (0.95, 0.15, 0.15), "rules": [
            ("z", 0.004, 0.020, 0.070, 0.086),       # behind the crown, forward of the mantle
            ("y", 0.266, 0.288, 0.360, 0.400),       # upper neck only — never the black bib/throat
            ("ny", -0.60, -0.35, 1.10, 1.20),        # skip the under-chin faces
        ]},
        # --- the shoulder / wing bar: on the flanks only, so the dorsal midline stays black
        {"name": "wing", "debug": (0.15, 0.9, 0.25), "rules": [
            ("ax", 0.032, 0.044, 0.090, 0.095),      # |x| — outer flank, both sides
            ("y", 0.172, 0.196, 0.252, 0.266),       # the TOP of the folded wing, not the flank
            ("z", -0.040, -0.020, 0.048, 0.064),     # anterior: the primaries stay black
        ]},
        # --- rump + tail base: white from the lower back into the proximal tail, black tip left
        {"name": "tailbase", "debug": (0.2, 0.35, 0.98), "rules": [
            ("z", -0.172, -0.156, -0.062, -0.044),   # black terminal band survives past -0.16
            ("y", -0.05, 0.0, 0.162, 0.184),         # under the back line, so the mantle stays black
        ]},
    ],
    # the paint: magpie white is a true bright white, faintly warm in sun
    "paint": (0.955, 0.950, 0.930),
}

SPECS = {"magpie": MAGPIE}


# ---------------------------------------------------------------- GLB read / write
def glb_read(path):
    b = open(path, "rb").read()
    total = struct.unpack("<I", b[8:12])[0]
    off, chunks = 12, []
    while off < total:
        clen, ctype = struct.unpack("<II", b[off:off + 8])
        chunks.append([ctype, bytearray(b[off + 8:off + 8 + clen])])
        off += 8 + clen
    return json.loads(chunks[0][1].decode("utf-8")), chunks


def glb_write(path, J, chunks):
    js = json.dumps(J, separators=(",", ":")).encode("utf-8")
    js += b" " * ((4 - len(js) % 4) % 4)
    bin_ = bytes(chunks[1][1])
    bin_ += b"\0" * ((4 - len(bin_) % 4) % 4)
    out = struct.pack("<III", 0x46546C67, 2, 12 + 8 + len(js) + 8 + len(bin_))
    out += struct.pack("<II", len(js), 0x4E4F534A) + js
    out += struct.pack("<II", len(bin_), 0x004E4942) + bin_
    open(path, "wb").write(out)
    return len(out)


def accessor(J, BIN, idx):
    a = J["accessors"][idx]
    bv = J["bufferViews"][a["bufferView"]]
    dt = {5126: "<f4", 5123: "<u2", 5125: "<u4", 5121: "u1"}[a["componentType"]]
    n = {"SCALAR": 1, "VEC2": 2, "VEC3": 3, "VEC4": 4}[a["type"]]
    arr = np.frombuffer(bytes(BIN), dtype=dt, count=a["count"] * n,
                        offset=bv["byteOffset"] + a.get("byteOffset", 0))
    return arr.reshape(a["count"], n) if n > 1 else arr


# ---------------------------------------------------------------- UV -> 3D rasterizer
def rasterize(P, N, UV, tris, size):
    """texel -> (position, normal, covered). Barycentric, top-left UV origin (glTF)."""
    pos = np.zeros((size, size, 3), np.float32)
    nrm = np.zeros((size, size, 3), np.float32)
    cov = np.zeros((size, size), bool)
    px = UV[:, 0] * size - 0.5
    py = UV[:, 1] * size - 0.5                      # v grows DOWN the image in glTF
    for t in tris:
        i0, i1, i2 = t
        x0, y0, x1, y1, x2, y2 = px[i0], py[i0], px[i1], py[i1], px[i2], py[i2]
        den = (y1 - y2) * (x0 - x2) + (x2 - x1) * (y0 - y2)
        if abs(den) < 1e-9:
            continue
        # 1-texel conservative pad so island interiors are fully covered
        cl = max(int(np.floor(min(x0, x1, x2))) - 1, 0)
        cr = min(int(np.ceil(max(x0, x1, x2))) + 1, size - 1)
        rt = max(int(np.floor(min(y0, y1, y2))) - 1, 0)
        rb = min(int(np.ceil(max(y0, y1, y2))) + 1, size - 1)
        if cr < cl or rb < rt:
            continue
        gx, gy = np.meshgrid(np.arange(cl, cr + 1), np.arange(rt, rb + 1))
        l0 = ((y1 - y2) * (gx - x2) + (x2 - x1) * (gy - y2)) / den
        l1 = ((y2 - y0) * (gx - x2) + (x0 - x2) * (gy - y2)) / den
        l2 = 1.0 - l0 - l1
        m = (l0 >= -0.25) & (l1 >= -0.25) & (l2 >= -0.25)       # -0.25 = the pad
        if not m.any():
            continue
        b = np.stack([np.clip(l0[m], 0, 1), np.clip(l1[m], 0, 1), np.clip(l2[m], 0, 1)], 1)
        b /= b.sum(1, keepdims=True)
        rr, cc = gy[m], gx[m]
        pos[rr, cc] = b @ P[[i0, i1, i2]]
        nrm[rr, cc] = b @ N[[i0, i1, i2]]
        cov[rr, cc] = True
    return pos, nrm, cov


def region_mask(spec, pos, nrm, cov):
    x, y, z = pos[..., 0], pos[..., 1], pos[..., 2]
    axes = {"x": x, "y": y, "z": z, "ax": np.abs(x),
            "nx": nrm[..., 0], "ny": nrm[..., 1], "nz": nrm[..., 2]}
    m = np.ones_like(y)
    for axis, lo_out, lo_in, hi_in, hi_out in spec["rules"]:
        m = m * window(axes[axis], lo_out, lo_in, hi_in, hi_out)
    return m * cov


def dilate(img, mask, cov, rounds=5):
    """Push painted colour outward into the bake's UV padding (uncovered texels next to painted
    ones), so bilinear/mip sampling at an island edge never pulls the old black through."""
    img = img.copy()
    done = (mask > 0.02)
    free = ~cov
    for _ in range(rounds):
        n_sum = np.zeros_like(img)
        n_cnt = np.zeros(img.shape[:2], np.float32)
        for dy, dx in ((1, 0), (-1, 0), (0, 1), (0, -1)):
            s = np.roll(np.roll(done.astype(np.float32), dy, 0), dx, 1)
            c = np.roll(np.roll(img, dy, 0), dx, 1)
            n_sum += c * s[..., None]
            n_cnt += s
        grow = free & (~done) & (n_cnt > 0)
        img[grow] = (n_sum[grow] / n_cnt[grow][..., None])
        done = done | grow
    return img


def main():
    which = sys.argv[1] if len(sys.argv) > 1 else "magpie"
    debug = "--debug" in sys.argv
    spec = SPECS[which]
    src = os.path.join(ROOT, "pipeline", "_normalized", spec["file"])
    out = sys.argv[sys.argv.index("--out") + 1] if "--out" in sys.argv else src
    J, chunks = glb_read(src)
    BIN = chunks[1][1]
    prim = J["meshes"][0]["primitives"][0]
    P = accessor(J, BIN, prim["attributes"]["POSITION"])
    N = accessor(J, BIN, prim["attributes"]["NORMAL"])
    UV = accessor(J, BIN, prim["attributes"]["TEXCOORD_0"])
    IDX = accessor(J, BIN, prim["indices"]).reshape(-1, 3)

    ibv = J["bufferViews"][J["images"][0]["bufferView"]]
    raw = bytes(BIN[ibv["byteOffset"]:ibv["byteOffset"] + ibv["byteLength"]])
    im = Image.open(io.BytesIO(raw)).convert("RGB")
    size = im.size[0]
    a = np.asarray(im).astype(np.float32) / 255.0
    print(f"{spec['file']}: {len(IDX)} tris · atlas {im.size} {J['images'][0]['mimeType']}")

    pos, nrm, cov = rasterize(P, N, UV, IDX, size)
    print(f"  rasterized: {cov.sum()} / {size*size} texels covered ({100*cov.mean():.1f}%)")

    out_img = a.copy()
    total = np.zeros(cov.shape, np.float32)
    for r in spec["regions"]:
        m = region_mask(r, pos, nrm, cov)
        total = np.maximum(total, m)
        col = np.array(r["debug"] if debug else spec["paint"], np.float32)
        if debug:
            out_img = out_img * (1 - m[..., None]) + col * m[..., None]
        else:
            # keep the bake's own light: its luminance, stretched into a white's shading range
            lum = a @ np.array([0.2126, 0.7152, 0.0722], np.float32)
            shade = 0.88 + 0.12 * smoothstep(0.0, 0.28, lum)
            painted = col * shade[..., None]
            out_img = out_img * (1 - m[..., None]) + painted * m[..., None]
        print(f"  {r['name']:9} {int((m>0.5).sum()):6d} texels solid, {int((m>0.02).sum()):6d} touched")
    out_img = dilate(np.clip(out_img, 0, 1), total, cov)

    png = Image.fromarray((np.clip(out_img, 0, 1) * 255 + 0.5).astype(np.uint8))
    if "--dump-atlas" in sys.argv:
        png.save(sys.argv[sys.argv.index("--dump-atlas") + 1])
    buf = io.BytesIO()
    png.save(buf, "WEBP", quality=88, method=6)
    new = buf.getvalue()

    # rewrite the image bufferView in place; it is the LAST view in the buffer (verified), so the
    # only offsets that move are its own length and the buffer's.
    tail_users = [bv for bv in J["bufferViews"] if bv["byteOffset"] > ibv["byteOffset"]]
    assert not tail_users, "image bufferView is not last — repack needed"
    del BIN[ibv["byteOffset"]:]
    BIN.extend(new)
    ibv["byteLength"] = len(new)
    J["buffers"][0]["byteLength"] = len(BIN)
    n = glb_write(out, J, chunks)
    print(f"  atlas {len(raw)} -> {len(new)} B webp · GLB -> {out} ({n} B)"
          + ("   [DEBUG COLOURS]" if debug else ""))


if __name__ == "__main__":
    main()