forgebench/code/ours/dump_seed.py

#!/usr/bin/env python
"""PASS 1 (REFERENCE ENV) — dump the ORACLE-VISIBILITY SS seed latent z1.

Three-pass SSFlow eval driver, pass 1 of 3:

    dump_seed.py      (reference env, production sam3d_objects)  -> z1 per object
    ssflow_coords.py  (TRAINING repo, vendored sam3d_objects)    -> coords per object
    finish_cache.py   (reference env)                            -> the stage-1 npz cache

WHY THREE PROCESSES: ``tools/val_stage1_ref.py`` deliberately pops PYTHONPATH so
that the PRODUCTION ``sam3d_objects`` (/lp-dev/jonghoon/mv-mesh/mv-sam3d) is
imported.  The trained ``SSFlowModel`` lives in the TRAINING repo's vendored
``sam3d_objects``.  Importing both in one interpreter is unsafe, so the pipeline
is split and the two sides only exchange .npz files.

WHAT THIS SCRIPT DOES — steps 1-4 of ``val_stage1_ref.stage1_and_seed()``,
BYTE-IDENTICAL:

  1. loads the input view PNGs ``EXP/inputs/<obj>_<tag>.png`` (existence +
     decodability check; they have NO numerical effect on z1 — in the reference
     they feed only the SS-generator conditioning, which is exactly the part
     pass 2 replaces).
  2. ``gt_vis, gt_full, free = baf.load_gt_synth(Path(exp), obj, views_tags, n=baf.VOX)``
  3. ``rng = np.random.default_rng(seed); occ = baf.build_input_grid(gt_vis, free, "zeros", rng)``
     -- the 64^3 ORACLE-VISIBILITY seed, union over the input views.
  4. ``z1 = baf.grid_to_shape_latent(ss_enc, occ)``  -- the SS-encoded partial
     occupancy == ``z_partial``.

Lines 2-4 are the reference block verbatim; ``noise="zeros"`` consumes nothing
from ``rng``, so z1 is a deterministic function of (exp, obj, views).

DELIBERATE DEVIATION (documented): the reference builds the WHOLE
``baf.Models()`` (production Inference pipeline + SLAT encoder + DINO).  z1 needs
only ``baf.load_ss_encoder()``, so that is all we build.  ``grid_to_shape_latent``
is a pure function of (ss_enc, occ), so the dumped z1 is bit-identical to the
reference's -- just ~100x cheaper to produce.  The da3 pointmaps
(``npz_{n}v/<obj>/da3_output.npz``) are NOT read: they only feed the SS
conditioning that pass 2 replaces.

Output: ``<out>/<obj>.npz`` with
    z1  (1,4096,8) float32   exactly as grid_to_shape_latent produced it
    occ (64,64,64) uint8     the partial-occupancy grid, for the record
    views (V,) str, seed int

Launch (re-execs itself into the reference env; caller sets CUDA_VISIBLE_DEVICES
to ONE gpu -- baf.DEVICE is hardcoded 'cuda'):
  CUDA_VISIBLE_DEVICES=7 python dump_seed.py \
      --exp /lp-dev/jonghoon/mv-mesh/exp_faithfulness/toys4k100_tex \
      --views 2 --out /lp-dev/jonghoon/ss_flow_data/eval_20260916/seed/toys4k100_tex/2v
"""
import os
import sys

ENV_PY = "/lp-dev/jonghoon/mv-mesh/envs/mv-sam3d/bin/python"
ENV = "/lp-dev/jonghoon/mv-mesh/envs/mv-sam3d"
MVMESH = "/lp-dev/jonghoon/mv-mesh"
METRICS_DIR = os.path.join(MVMESH, "metrics")

# --- re-exec preamble: VERBATIM from tools/val_stage1_ref.py -----------------
if not os.environ.get("_VAL_STAGE1_INENV"):
    env = dict(os.environ)
    env["_VAL_STAGE1_INENV"] = "1"
    env["_APPFORCE_SAM3D_INENV"] = "1"       # batch_appforce must not re-exec
    env.pop("PYTHONPATH", None)              # never the training repo's sam3d_objects
    env["PATH"] = os.path.join(ENV, "bin") + os.pathsep + env.get("PATH", "")
    assert env.get("CUDA_VISIBLE_DEVICES"), "set CUDA_VISIBLE_DEVICES to one GPU"
    os.execve(ENV_PY, [ENV_PY, os.path.abspath(__file__)] + sys.argv[1:], env)

os.environ["_APPFORCE_SAM3D_INENV"] = "1"
sys.path.insert(0, METRICS_DIR)
import batch_appforce_sam3d as baf  # noqa: E402  (chdir REPO, sys.path, env defaults)

import argparse  # noqa: E402
import json      # noqa: E402
import time      # noqa: E402
import traceback  # noqa: E402
from pathlib import Path  # noqa: E402

import numpy as np  # noqa: E402
import torch        # noqa: E402
from PIL import Image  # noqa: E402

# --- VIEW TAG TABLE: verbatim from tools/val_stage1_ref.py ------------------
VIEW_TAGS = ["front", "side", "back", "oside", "top", "bottom", "top2", "bottom2"]


def views_tags_for(n: int):
    assert 1 <= n <= len(VIEW_TAGS), f"--views {n} unsupported (max {len(VIEW_TAGS)})"
    return VIEW_TAGS[:n]


def _objects(args):
    """Verbatim from tools/val_stage1_ref.py::_objects."""
    if args.objects_file:
        objs = [l.strip() for l in open(args.objects_file) if l.strip()]
    elif args.objects:
        objs = list(args.objects)
    else:
        data = json.loads(Path(os.path.join(args.exp, "selection.json")).read_text())
        sels = data["selections"] if isinstance(data, dict) else data
        objs = [s["object"] for s in sels]
    if args.limit:
        objs = objs[: args.limit]
    if args.nshards > 1:
        objs = objs[args.shard::args.nshards]
    return objs


def seed_for_object(ss_enc, obj, exp, inputs_dir, views_tags, seed):
    """Steps 1-4 of val_stage1_ref.stage1_and_seed, unchanged."""
    # ---- step 1 (verbatim; existence/decodability check only) --------------
    for t in views_tags:
        p = os.path.join(inputs_dir, f"{obj}_{t}.png")
        if not os.path.isfile(p):
            raise FileNotFoundError(p)
        np.array(Image.open(p).convert("RGBA"))

    # ---- steps 2-4 (verbatim) ----------------------------------------------
    gt_vis, gt_full, free = baf.load_gt_synth(Path(exp), obj, views_tags, n=baf.VOX)
    rng = np.random.default_rng(seed)
    occ = baf.build_input_grid(gt_vis, free, "zeros", rng)
    z1 = baf.grid_to_shape_latent(ss_enc, occ)

    return dict(z1=z1.detach().float().cpu().numpy(),      # (1,4096,8) float32
                occ=(occ > 0.5).astype(np.uint8))          # (64,64,64) uint8


def main():
    ap = argparse.ArgumentParser(description=__doc__,
                                 formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--exp", required=True,
                    help="EXPDIR (renders/, inputs/, selection.json)")
    ap.add_argument("--views", type=int, default=2,
                    help="number of condition views (1..8); tags = VIEW_TAGS[:views]")
    ap.add_argument("--out", required=True, help="seed dump dir (<obj>.npz each)")
    ap.add_argument("--seed", type=int, default=42)
    ap.add_argument("--objects", nargs="+", default=None)
    ap.add_argument("--objects-file", default=None)
    ap.add_argument("--limit", type=int, default=0)
    ap.add_argument("--shard", type=int, default=0)
    ap.add_argument("--nshards", type=int, default=1)
    ap.add_argument("--force", action="store_true")
    args = ap.parse_args()

    exp = os.path.abspath(args.exp)
    inputs_dir = os.path.join(exp, "inputs")
    views_tags = views_tags_for(args.views)
    out = Path(args.out)
    out.mkdir(parents=True, exist_ok=True)
    objs = _objects(args)
    print(f"[seed] exp={exp} views={args.views} n_obj={len(objs)} "
          f"shard={args.shard}/{args.nshards} seed={args.seed} "
          f"CUDA={os.environ.get('CUDA_VISIBLE_DEVICES')}", flush=True)
    t0 = time.time()
    ss_enc = baf.load_ss_encoder()
    print(f"[seed] ss_encoder ready {time.time() - t0:.1f}s", flush=True)

    n_ok = n_fail = n_skip = 0
    for i, obj in enumerate(objs, 1):
        f = out / f"{obj}.npz"
        if f.is_file() and not args.force:
            n_skip += 1
            continue
        t1 = time.time()
        try:
            d = seed_for_object(ss_enc, obj, exp, inputs_dir, views_tags, args.seed)
            tmp = str(f) + ".tmp.npz"
            np.savez_compressed(tmp, **d, views=np.array(views_tags), seed=args.seed)
            os.replace(tmp, f)
            n_ok += 1
            print(f"[seed {i}/{len(objs)}] OK {obj} occ={int(d['occ'].sum())} "
                  f"z1={d['z1'].shape} {time.time() - t1:.1f}s", flush=True)
        except Exception:
            n_fail += 1
            traceback.print_exc()
            print(f"[seed {i}/{len(objs)}] FAIL {obj}", flush=True)
        torch.cuda.empty_cache()
    print(f"DUMPSEED DONE ok={n_ok} fail={n_fail} skip={n_skip} total={len(objs)}",
          flush=True)


if __name__ == "__main__":
    main()