| |
| """Build a 3DEditFormer eval adapter for Edit3D-Bench.""" |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import json |
| import os |
| from pathlib import Path |
| from typing import Any |
|
|
|
|
| DEFAULT_SS_NAME = "ss_enc_conv3d_16l8_fp16" |
| DEFAULT_LATENT_NAME = "dinov2_vitl14_reg_slat_enc_swin8_B_64l8_fp16" |
|
|
|
|
| def task_id(dataset: str, source_model: str, prompt_key: str | None = None) -> str: |
| base = f"{dataset}__{source_model}" |
| return f"{base}__{prompt_key}" if prompt_key else base |
|
|
|
|
| def rel(path: Path, root: Path) -> str: |
| return path.relative_to(root).as_posix() |
|
|
|
|
| def replace_link(dst: Path, src: Path, *, copy_images: 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_images and src.suffix.lower() in {".png", ".jpg", ".jpeg"}: |
| import shutil |
|
|
| shutil.copy2(src, dst) |
| else: |
| os.symlink(src, dst) |
|
|
|
|
| def load_metadata(bench_root: Path) -> list[dict[str, Any]]: |
| metadata_path = bench_root / "metadata.json" |
| with metadata_path.open("r", encoding="utf-8") as f: |
| metadata = json.load(f) |
| if not isinstance(metadata, list): |
| raise TypeError(f"Expected list metadata in {metadata_path}") |
| return metadata |
|
|
|
|
| def choose_ori_image(bench_task_dir: Path, preprocess_render_dir: Path) -> Path: |
| candidates = [ |
| bench_task_dir / "2d_render.png", |
| preprocess_render_dir / "000.png", |
| bench_task_dir.parent / "source_model" / "render" / "render_0000.png", |
| ] |
| for candidate in candidates: |
| if candidate.exists(): |
| return candidate |
| return candidates[0] |
|
|
|
|
| def build_adapter(args: argparse.Namespace) -> dict[str, Any]: |
| bench_root = Path(args.bench_root).resolve() |
| preprocess_root = Path(args.preprocess_root).resolve() |
| output_root = Path(args.output_root).resolve() |
| flux_root = output_root / "flux_edit" |
|
|
| metadata = load_metadata(bench_root) |
| data_info: dict[str, dict[str, dict[str, str]]] = {"flux_edit": {}} |
| select: dict[str, list[str]] = {"flux_edit": []} |
| task_map: dict[str, dict[str, Any]] = {} |
| missing: list[str] = [] |
|
|
| ss_root = preprocess_root / "ss_latents" / args.ss_latent_name |
| latent_root = preprocess_root / "latents" / args.latent_name |
| render_root = preprocess_root / "renders" |
|
|
| for item in metadata: |
| dataset = item["dataset"] |
| source_model = item["source_model"] |
| source_id = task_id(dataset, source_model) |
| ss_path = ss_root / f"{source_id}.npz" |
| latent_path = latent_root / f"{source_id}.npz" |
| preprocess_render_dir = render_root / source_id |
|
|
| for prompt_idx in range(1, 4): |
| prompt_key = f"prompt_{prompt_idx}" |
| if prompt_key not in item: |
| continue |
|
|
| key = task_id(dataset, source_model, prompt_key) |
| task_dir = flux_root / key |
| bench_task_dir = bench_root / dataset / source_model / prompt_key |
| edit_img = bench_task_dir / "2d_edit.png" |
| ori_img = choose_ori_image(bench_task_dir, preprocess_render_dir) |
|
|
| required = { |
| "ss_latent": ss_path, |
| "latent": latent_path, |
| "ori_img": ori_img, |
| "edit_img": edit_img, |
| "source_model": bench_root / dataset / source_model / "source_model" / "model.glb", |
| "edit_region": bench_task_dir / "3d_edit_region.glb", |
| } |
| for name, req_path in required.items(): |
| if not req_path.exists(): |
| missing.append(f"{key}: missing {name}: {req_path}") |
|
|
| replace_link(task_dir / "ori_ss_latents.npz", ss_path) |
| replace_link(task_dir / "ori_latents.npz", latent_path) |
| replace_link(task_dir / "ori_img.png", ori_img, copy_images=args.copy_images) |
| replace_link(task_dir / "edit_img.png", edit_img, copy_images=args.copy_images) |
|
|
| data_info["flux_edit"][key] = { |
| "ori_ss_latents_path": rel(task_dir / "ori_ss_latents.npz", flux_root), |
| "ori_latents_path": rel(task_dir / "ori_latents.npz", flux_root), |
| "ori_img_path": rel(task_dir / "ori_img.png", flux_root), |
| "edit_img_path": rel(task_dir / "edit_img.png", flux_root), |
| } |
| select["flux_edit"].append(key) |
| task_map[key] = { |
| "dataset": dataset, |
| "source_model": source_model, |
| "prompt_key": prompt_key, |
| "prompt": item[prompt_key], |
| "source_prompt": item.get("source_prompt", ""), |
| "source_id": source_id, |
| } |
|
|
| output_root.mkdir(parents=True, exist_ok=True) |
| (output_root / "dataset_info.json").write_text( |
| json.dumps(data_info, ensure_ascii=False, indent=2) + "\n", encoding="utf-8" |
| ) |
| (output_root / "test_data_info.json").write_text( |
| json.dumps(select, ensure_ascii=False, indent=2) + "\n", encoding="utf-8" |
| ) |
| smoke_select = {"flux_edit": select["flux_edit"][: args.smoke_count]} |
| (output_root / "test_data_info_smoke.json").write_text( |
| json.dumps(smoke_select, ensure_ascii=False, indent=2) + "\n", encoding="utf-8" |
| ) |
| (output_root / "edit3d_task_map.json").write_text( |
| json.dumps(task_map, ensure_ascii=False, indent=2) + "\n", encoding="utf-8" |
| ) |
|
|
| summary = { |
| "output_root": str(output_root), |
| "num_sources": len(metadata), |
| "num_tasks": len(select["flux_edit"]), |
| "smoke_tasks": len(smoke_select["flux_edit"]), |
| "missing": missing, |
| } |
| (output_root / "adapter_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("--bench_root", required=True) |
| parser.add_argument("--preprocess_root", required=True) |
| parser.add_argument("--output_root", required=True) |
| parser.add_argument("--ss_latent_name", default=DEFAULT_SS_NAME) |
| parser.add_argument("--latent_name", default=DEFAULT_LATENT_NAME) |
| parser.add_argument("--smoke_count", type=int, default=3) |
| parser.add_argument("--copy_images", action="store_true") |
| parser.add_argument("--strict", action="store_true") |
| args = parser.parse_args() |
|
|
| summary = build_adapter(args) |
| print(json.dumps(summary, ensure_ascii=False, indent=2)) |
| if args.strict and summary["missing"]: |
| raise SystemExit(1) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
