File size: 3,819 Bytes
2274e38 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 | #!/usr/bin/env python3
"""Export 3DEditFormer Edit3D-Bench predictions to official evaluator layout."""
from __future__ import annotations
import argparse
import json
import os
import shutil
from pathlib import Path
from typing import Any
def replace_link(dst: Path, src: Path, *, copy: bool = False) -> None:
if dst.exists() or dst.is_symlink():
if dst.is_symlink() and Path(os.readlink(dst)) == src:
return
dst.unlink()
dst.parent.mkdir(parents=True, exist_ok=True)
if copy:
shutil.copy2(src, dst)
else:
os.symlink(src, dst)
def render_id(path: Path) -> str | None:
stem = path.stem
if stem.startswith("render_"):
suffix = stem.removeprefix("render_")
else:
suffix = stem
if suffix.isdigit():
return f"{int(suffix):04d}"
return None
def load_task_map(adapter_root: Path) -> dict[str, dict[str, Any]]:
path = adapter_root / "edit3d_task_map.json"
with path.open("r", encoding="utf-8") as f:
task_map = json.load(f)
if not isinstance(task_map, dict):
raise TypeError(f"Expected dict task map in {path}")
return task_map
def export_predictions(args: argparse.Namespace) -> dict[str, Any]:
adapter_root = Path(args.adapter_root).resolve()
eval_results_dir = Path(args.eval_results_dir).resolve()
pred_root = Path(args.pred_root).resolve()
task_map = load_task_map(adapter_root)
exported_models = 0
exported_images = 0
missing: list[str] = []
for key, meta in task_map.items():
src_task_dir = eval_results_dir / "flux_edit" / key
src_mesh = src_task_dir / "edit.glb"
src_render_dir = src_task_dir / "render"
dst_task_dir = pred_root / meta["dataset"] / meta["source_model"] / meta["prompt_key"]
dst_mesh = dst_task_dir / "edit.glb"
dst_image_dir = dst_task_dir / "images"
if src_mesh.exists():
replace_link(dst_mesh, src_mesh, copy=args.copy)
exported_models += 1
else:
missing.append(f"{key}: missing mesh: {src_mesh}")
if not src_render_dir.exists():
missing.append(f"{key}: missing render dir: {src_render_dir}")
continue
render_files = sorted(
p for p in src_render_dir.glob("*.png") if render_id(p) is not None
)
if not render_files:
missing.append(f"{key}: no render pngs in {src_render_dir}")
continue
for src_img in render_files:
img_id = render_id(src_img)
if img_id is None:
continue
dst_img = dst_image_dir / f"render_{img_id}.png"
replace_link(dst_img, src_img, copy=args.copy)
exported_images += 1
summary = {
"pred_root": str(pred_root),
"num_tasks": len(task_map),
"exported_models": exported_models,
"exported_images": exported_images,
"missing": missing,
}
pred_root.mkdir(parents=True, exist_ok=True)
(pred_root / "export_summary.json").write_text(
json.dumps(summary, ensure_ascii=False, indent=2) + "\n", encoding="utf-8"
)
return summary
def main() -> None:
parser = argparse.ArgumentParser()
parser.add_argument("--adapter_root", required=True)
parser.add_argument("--eval_results_dir", required=True)
parser.add_argument("--pred_root", required=True)
parser.add_argument("--copy", action="store_true", help="Copy files instead of symlinking")
parser.add_argument("--strict", action="store_true")
args = parser.parse_args()
summary = export_predictions(args)
print(json.dumps(summary, ensure_ascii=False, indent=2))
if args.strict and summary["missing"]:
raise SystemExit(1)
if __name__ == "__main__":
main()
|