| |
| """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() |
|
|
