| |
| """Batch Cupid3D driver: load the pipeline ONCE, process an object list. |
| |
| Usage: |
| python batch_cupid.py --selection SEL.json --inputs DIR --out OUTDIR --views 1 [--seed 42] |
| |
| SEL.json: {"selections": [{"object": NAME, ...}, ...]} |
| Images: DIR/<object>_front.png RGBA crops (Cupid uses the existing alpha as-is; |
| _has_alpha() -> True so NO rembg re-segmentation is triggered). |
| Output: OUTDIR/<object>.glb -- Cupid's mesh in its OWN canonical model frame, |
| exported exactly like Cupid's save_mesh(): |
| postprocessing_utils.to_glb(gaussian, mesh, simplify=0.95, texture_size=1024). |
| The estimated camera pose is intentionally NOT applied: downstream |
| metrics/align_baselines.py registers each prediction to the GT canonical |
| frame (best-of-24 cube rotation + rigid ICP), identical to RVG/Pixal3D/ |
| Amodal3R/Hunyuan. So the raw model-frame mesh is the fair input. |
| |
| Cupid only exposes single-image reconstruction (pipeline.run takes ONE image), |
| so only --views 1 is supported here. |
| |
| Idempotent: existing OUTDIR/<object>.glb is skipped. Per-object try/except. |
| Re-execs itself inside the cupid conda env if not already there. |
| """ |
| import os |
| import sys |
|
|
| ENV = "/lp-dev/jonghoon/mv-mesh/baselines/cupid/env" |
| ENV_PY = os.path.join(ENV, "bin", "python") |
| REPO = "/lp-dev/jonghoon/mv-mesh/baselines/cupid/repo" |
| HF = "/lp-dev/jonghoon/mv-mesh/hf_cache" |
|
|
| _ENV_VARS = { |
| "PYTHONUNBUFFERED": "1", |
| "OMP_NUM_THREADS": "4", |
| "MKL_NUM_THREADS": "4", |
| "SPCONV_ALGO": "native", |
| "ATTN_BACKEND": "flash_attn", |
| "MKL_THREADING_LAYER": "GNU", |
| "MKL_SERVICE_FORCE_INTEL": "0", |
| "HF_HOME": HF, |
| "HUGGINGFACE_HUB_CACHE": HF, |
| "HF_HUB_CACHE": HF, |
| "TORCH_HOME": "/lp-dev/jonghoon/mv-mesh/torch_hub", |
| "CUDA_HOME": "/usr/local/cuda-12.8", |
| } |
|
|
|
|
| def _reexec_in_env(): |
| env = dict(os.environ) |
| for k, v in _ENV_VARS.items(): |
| env[k] = v |
| for k in ("OMP_NUM_THREADS", "MKL_NUM_THREADS"): |
| if os.environ.get(k): |
| env[k] = os.environ[k] |
| env["CONDA_PREFIX"] = ENV |
| env["PATH"] = os.path.join(ENV, "bin") + os.pathsep + "/usr/local/cuda-12.8/bin" + os.pathsep + env.get("PATH", "") |
| if not env.get("CUDA_VISIBLE_DEVICES"): |
| try: |
| import subprocess |
| out = subprocess.check_output( |
| ["nvidia-smi", "--query-gpu=index,memory.free", |
| "--format=csv,noheader,nounits"], text=True) |
| best = max((l.split(",") for l in out.strip().splitlines()), |
| key=lambda r: int(r[1])) |
| env["CUDA_VISIBLE_DEVICES"] = best[0].strip() |
| except Exception as e: |
| print(f"[batch] (non-fatal) GPU autopick failed: {e}", flush=True) |
| env["_CUPID_BATCH_INENV"] = "1" |
| print(f"[batch] re-exec in {ENV_PY} " |
| f"(CUDA_VISIBLE_DEVICES={env.get('CUDA_VISIBLE_DEVICES')})", flush=True) |
| os.execve(ENV_PY, [ENV_PY, os.path.abspath(__file__)] + sys.argv[1:], env) |
|
|
|
|
| if not os.environ.get("_CUPID_BATCH_INENV"): |
| if not os.path.isfile(ENV_PY): |
| sys.exit(f"ERROR: env python not found: {ENV_PY}") |
| _reexec_in_env() |
|
|
| for _k, _v in _ENV_VARS.items(): |
| os.environ.setdefault(_k, _v) |
|
|
| import argparse |
| import json |
| import time |
| import traceback |
|
|
| os.chdir(REPO) |
| sys.path.insert(0, REPO) |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--selection", required=True) |
| ap.add_argument("--inputs", required=True) |
| ap.add_argument("--out", required=True) |
| ap.add_argument("--views", type=int, default=1, choices=[1]) |
| ap.add_argument("--seed", type=int, default=42) |
| ap.add_argument("--limit", type=int, default=None) |
| ap.add_argument("--simplify", type=float, default=0.95) |
| ap.add_argument("--texture-size", type=int, default=1024) |
| ap.add_argument("--shard", type=int, default=0) |
| ap.add_argument("--nshards", type=int, default=1) |
| ap.add_argument("--reverse", action="store_true") |
| args = ap.parse_args() |
|
|
| with open(args.selection) as f: |
| sel = json.load(f)["selections"] |
| objects = [s["object"] for s in sel] |
| if args.limit: |
| objects = objects[: args.limit] |
| if args.nshards > 1: |
| objects = objects[args.shard::args.nshards] |
| if args.reverse: |
| objects = objects[::-1] |
| outdir = os.path.abspath(args.out) |
| os.makedirs(outdir, exist_ok=True) |
| inputs_dir = os.path.join(outdir, "inputs") |
| os.makedirs(inputs_dir, exist_ok=True) |
|
|
| jobs = [] |
| for name in objects: |
| out = os.path.join(outdir, f"{name}.glb") |
| img = os.path.join(args.inputs, f"{name}_front.png") |
| jobs.append((name, img, out)) |
|
|
| todo = [j for j in jobs if not os.path.isfile(j[2])] |
| print(f"[batch] {len(jobs)} object(s), {len(jobs) - len(todo)} already done, " |
| f"{len(todo)} to run | views={args.views} seed={args.seed}", flush=True) |
|
|
| import torch |
| from PIL import Image |
| from cupid.pipelines import Cupid3DPipeline |
| from cupid.utils import sample_utils |
| from cupid.utils import postprocessing_utils |
|
|
| n_ok = n_fail = n_skip = 0 |
| pipeline = None |
|
|
| for i, (name, img_path, out) in enumerate(jobs, 1): |
| t0 = time.time() |
| try: |
| if os.path.isfile(out): |
| n_skip += 1 |
| print(f"[{i}/{len(jobs)}] {name} SKIP (exists)", flush=True) |
| continue |
| if not os.path.isfile(img_path): |
| raise FileNotFoundError(img_path) |
|
|
| if pipeline is None: |
| tl = time.time() |
| pipeline = Cupid3DPipeline.from_pretrained("hbb1/Cupid") |
| pipeline.cuda() |
| print(f"[batch] pipeline ready in {time.time()-tl:.1f}s", flush=True) |
| t0 = time.time() |
|
|
| image = sample_utils.load_image(img_path) |
| outputs = pipeline.run(image, seed=args.seed, |
| formats=["mesh", "gaussian"]) |
| |
| try: |
| proc = pipeline.crop_image(pipeline.preprocess_image(image)) |
| proc.image.save(os.path.join(inputs_dir, f"FED_{name}_front.png")) |
| except Exception: |
| pass |
|
|
| glb = postprocessing_utils.to_glb( |
| outputs["gaussian"][0], outputs["mesh"][0], |
| simplify=args.simplify, texture_size=args.texture_size, |
| verbose=False, |
| ) |
| tmp = out + ".tmp.glb" |
| glb.export(tmp) |
| if not os.path.isfile(tmp): |
| raise RuntimeError("export produced no file") |
| os.replace(tmp, out) |
| del outputs, glb, image |
| n_ok += 1 |
| print(f"[{i}/{len(jobs)}] {name} OK {time.time()-t0:.1f}s -> {out}", flush=True) |
| except Exception: |
| n_fail += 1 |
| traceback.print_exc() |
| print(f"[{i}/{len(jobs)}] {name} FAIL {time.time()-t0:.1f}s", flush=True) |
| finally: |
| try: |
| torch.cuda.empty_cache() |
| except Exception: |
| pass |
|
|
| print(f"[batch] DONE ok={n_ok} fail={n_fail} skip={n_skip}", flush=True) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
