forgebench/code/build_forgebench_1v.py

#!/usr/bin/env python3
"""Build FORGE3DBench 1-view eval exp.
For every scene / every spawned object: compute per-view visibility =
modal_area(mask/) / amodal_area(mask_amodal/). Keep object if >=1 view has
vis>=0.40; pick ONE qualifying view at RANDOM (global seed 0). Writes standard
exp (inputs/<uid>_front.png RGBA crop, renders/<uid>_canon.glb, renders/<uid>_front.npz,
npz_1v/<uid>/da3_output.npz). Conventions copied from .debug_psl_qual/build_psl_qual.py.
"""
import os, sys, json, random, re, numpy as np, trimesh
from PIL import Image
from pathlib import Path
from concurrent.futures import ProcessPoolExecutor

ROOT = "/data/psl_benchmark/MVScenes_Benchmark"
OUT = Path("/lp-dev/jonghoon/mv-mesh/exp_faithfulness/forgebench_1v")
MANIFEST = Path("/lp-dev/jonghoon/mv-mesh/.debug/forgebench_eval/manifest.json")
SIZE = 37
TARGET_SPAN = 0.9
VIS_THR = 0.40
SEED = 0
NCAM = 24
ANCHORS = [(e, a) for e in ("Factory", "Park")
           for a in ("Anchor_Ground", "Anchor_Table", "Anchor_Robot_Table")]
R_FIX = np.array([[1, 0, 0, 0], [0, 0, -1, 0], [0, 1, 0, 0], [0, 0, 0, 1]], float)


def uid_for(env, anch, scdir, obj_id):
    a = anch.replace("Anchor_", "")
    sc = scdir.replace("scene_", "s")
    safe = re.sub(r"[^0-9a-zA-Z]+", "", obj_id)
    return f"{env[:2]}{a[:2]}_{sc}_{safe}"


def crop_bbox(mask):
    ys, xs = np.nonzero(mask)
    y0, y1, x0, x1 = ys.min(), ys.max() + 1, xs.min(), xs.max() + 1
    mh, mw = int((y1 - y0) * 0.15) + 8, int((x1 - x0) * 0.15) + 8
    h, w = mask.shape
    y0, y1 = max(0, y0 - mh), min(h, y1 + mh)
    x0, x1 = max(0, x0 - mw), min(w, x1 + mw)
    side = max(y1 - y0, x1 - x0)
    cy, cx = (y0 + y1) // 2, (x0 + x1) // 2
    y0, y1 = max(0, cy - side // 2), min(h, cy + (side + 1) // 2)
    x0, x1 = max(0, cx - side // 2), min(w, cx + (side + 1) // 2)
    return int(y0), int(y1), int(x0), int(x1)


def pointmap_from_depth(depth_mm, K, bbox):
    y0, y1, x0, x1 = bbox
    dc = depth_mm[y0:y1, x0:x1].astype(np.float64) / 1000.0
    H, W = dc.shape
    v, u = np.meshgrid(np.arange(H), np.arange(W), indexing="ij")
    x = (u + x0 - K["cx"]) * dc / K["fx"]
    y = (v + y0 - K["cy"]) * dc / K["fy"]
    pm = np.stack([x, y, dc], -1).astype(np.float32)
    pm[dc == 0] = np.nan
    ri = np.clip((np.arange(SIZE) + 0.5) * H / SIZE, 0, H - 1).astype(int)
    ci = np.clip((np.arange(SIZE) + 0.5) * W / SIZE, 0, W - 1).astype(int)
    return pm[ri][:, ci].transpose(2, 0, 1)


def process_scene(args):
    env, anch, scdir = args
    p = f"{ROOT}/{env}/{anch}/scenes/{scdir}"
    try:
        cams = json.load(open(f"{p}/cameras.json"))["cameras"]
        scene = json.load(open(f"{p}/scene.json"))
        poses = {o["obj_id"]: o for o in json.load(open(f"{p}/poses.json"))["objects"]}
    except Exception as e:
        return [], [f"SCENE_LOAD_FAIL {env}/{anch}/{scdir}: {e}"]
    objs = [o for o in scene["objects"] if o.get("role") == "object" and o.get("spawned")]
    if not objs:
        return [], []
    ncam = min(len(cams), NCAM)
    # load modal masks (RGB) once per cam
    modal = {}
    for ci in range(ncam):
        mp = f"{p}/mask/cam{ci:02d}.png"
        if os.path.exists(mp):
            modal[ci] = np.array(Image.open(mp).convert("RGB")).astype(np.int16)
    rng = random.Random(SEED)  # base; per-object reseeded for determinism
    recs = []
    logs = []
    for o in objs:
        oid = o["obj_id"]; asset = o["asset"]; mc = np.array(o["mask_color"], dtype=np.int16)
        vis = np.zeros(ncam)
        modal_area = np.zeros(ncam, int)
        for ci in range(ncam):
            if ci not in modal:
                continue
            sel = (np.abs(modal[ci] - mc).sum(2) <= 8)
            ma = int(sel.sum())
            modal_area[ci] = ma
            amp = f"{p}/mask_amodal/cam{ci:02d}/{oid}.png"
            if not os.path.exists(amp):
                continue
            am = np.array(Image.open(amp).convert("L"))
            aa = int((am > 0).sum())
            if aa > 0:
                vis[ci] = ma / aa
        qual = [ci for ci in range(ncam) if vis[ci] >= VIS_THR]
        if not qual:
            continue
        # deterministic per-object random choice among qualifying views
        r = random.Random(f"{env}/{anch}/{scdir}/{oid}/{SEED}")
        chosen = r.choice(sorted(qual))
        uid = uid_for(env, anch, scdir, oid)
        try:
            rec = build_object(env, anch, scdir, p, cams, o, poses, chosen, uid)
            rec.update(dict(visibility=round(float(vis[chosen]), 4),
                            n_qual=len(qual), modal_area=int(modal_area[chosen])))
            recs.append(rec)
        except Exception as e:
            import traceback
            logs.append(f"BUILD_FAIL {uid}: {e}\n{traceback.format_exc()}")
    return recs, logs


def build_object(env, anch, scdir, p, cams, objrec, poses, ci, uid):
    asset = objrec["asset"]; oid = objrec["obj_id"]
    mc = np.array(objrec["mask_color"], dtype=np.int16)
    Two = np.array(poses[oid]["T_world_obj"])
    m = trimesh.load(f"{ROOT}/{env}/{anch}/objects/{asset}.glb", force="mesh")
    m.apply_transform(R_FIX)
    vw = (Two @ np.c_[m.vertices, np.ones(len(m.vertices))].T).T[:, :3]
    lo, hi = vw.min(0), vw.max(0)
    center = 0.5 * (lo + hi); extent = float((hi - lo).max())
    scale_r = extent / TARGET_SPAN
    ren = OUT / "renders"; inp = OUT / "inputs"
    ren.mkdir(parents=True, exist_ok=True); inp.mkdir(parents=True, exist_ok=True)
    mc2 = m.copy(); mc2.vertices = (vw - center) / scale_r
    mc2.export(ren / f"{uid}_canon.glb")
    cam = cams[ci]; Kd = cam["intrinsics"]
    fx, fy, cx, cy = Kd["fx"], Kd["fy"], Kd["cx"], Kd["cy"]
    c2w = np.array(cam["extrinsic_c2w"]).copy()
    c2w[:3, 3] = (c2w[:3, 3] - center) / scale_r
    modal = np.array(Image.open(f"{p}/mask/cam{ci:02d}.png").convert("RGB")).astype(np.int16)
    sel = (np.abs(modal - mc).sum(2) <= 8)
    raw = np.array(Image.open(f"{p}/depth/cam{ci:02d}.png")).astype(np.float64)
    sel = sel & (raw > 0)
    dmm = np.zeros(raw.shape, np.float64); dmm[sel] = raw[sel] / scale_r
    dmm = np.clip(np.rint(dmm), 0, 65535).astype(np.uint16)
    bb = crop_bbox(sel); y0, y1, x0, x1 = bb
    rgb = np.array(Image.open(f"{p}/rgb/cam{ci:02d}.png").convert("RGB"))
    alpha = np.where(sel, 255, 0).astype(np.uint8)
    rgba = np.dstack([rgb, alpha])[y0:y1, x0:x1]
    Image.fromarray(rgba, "RGBA").save(inp / f"{uid}_front.png")
    np.savez_compressed(ren / f"{uid}_front.npz", depth_mm=dmm, fx=float(fx), fy=float(fy),
                        cx=float(cx), cy=float(cy), c2w_cv=c2w, bbox=np.array(bb, np.int64),
                        res=int(raw.shape[0]), up_axis="y")
    d = np.load(ren / f"{uid}_front.npz")
    Kv = {k: float(d[k]) for k in ("fx", "fy", "cx", "cy")}
    pm = pointmap_from_depth(d["depth_mm"], Kv, tuple(d["bbox"]))
    dd = OUT / "npz_1v" / uid; dd.mkdir(parents=True, exist_ok=True)
    np.savez(dd / "da3_output.npz", pointmaps_sam3d=pm[None],
             image_files=np.array(["v1.png"]))
    return dict(object=uid, env=env, anchor=anch, scene=scdir, asset=asset, obj_id=oid,
                mask_color=objrec["mask_color"], scale_r=float(scale_r),
                center=center.tolist(), extent=extent, cam=int(ci),
                src_scene=p, front=int(ci))


def main():
    OUT.mkdir(parents=True, exist_ok=True)
    jobs = []
    for env, anch in ANCHORS:
        sd = f"{ROOT}/{env}/{anch}/scenes"
        if not os.path.isdir(sd):
            continue
        for scdir in sorted(os.listdir(sd)):
            if scdir.startswith("scene_"):
                jobs.append((env, anch, scdir))
    print(f"[build] {len(jobs)} scenes, workers=32, seed={SEED}, vis_thr={VIS_THR}", flush=True)
    sels = []; alllogs = []
    with ProcessPoolExecutor(max_workers=32) as ex:
        for i, (recs, logs) in enumerate(ex.map(process_scene, jobs)):
            sels.extend(recs); alllogs.extend(logs)
            if (i + 1) % 12 == 0:
                print(f"[build] {i+1}/{len(jobs)} scenes done, objs so far={len(sels)}", flush=True)
    for lg in alllogs:
        print("LOG:", lg, flush=True)
    sels.sort(key=lambda r: r["object"])
    json.dump({"source": "forgebench_mvscenes_1v", "seed": SEED, "vis_thr": VIS_THR,
               "n": len(sels), "selections": sels},
              open(OUT / "selection.json", "w"), indent=1)
    json.dump({"seed": SEED, "vis_thr": VIS_THR, "n_objects": len(sels),
               "n_scenes": len(jobs), "n_build_fail": len(alllogs),
               "selections": sels}, open(MANIFEST, "w"), indent=1)
    print(f"[build] DONE n_objects={len(sels)} -> {OUT/'selection.json'}", flush=True)
    print(f"[build] manifest -> {MANIFEST}", flush=True)


if __name__ == "__main__":
    main()