#!/usr/bin/env python3 """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()