#!/usr/bin/env python3 """Prepare Edit3D-Bench source models for 3DEditFormer/TRELLIS latent encoding. This script intentionally avoids copying the benchmark. It creates the minimal 3DEditFormer preprocessing tree expected by dataset_toolkits/extract_feature.py, encode_ss_latent.py, and encode_latent.py. """ import argparse import json import math import os import subprocess from concurrent.futures import ThreadPoolExecutor, as_completed from pathlib import Path import numpy as np import pandas as pd def sphere_hammersley_sequence(i: int, n: int, offset=(0.0, 0.0)): """Return yaw/pitch following the 3DEditFormer dataset toolkit convention.""" def radical_inverse(bits): bits = (bits << 16) | (bits >> 16) bits = ((bits & 0x55555555) << 1) | ((bits & 0xAAAAAAAA) >> 1) bits = ((bits & 0x33333333) << 2) | ((bits & 0xCCCCCCCC) >> 2) bits = ((bits & 0x0F0F0F0F) << 4) | ((bits & 0xF0F0F0F0) >> 4) bits = ((bits & 0x00FF00FF) << 8) | ((bits & 0xFF00FF00) >> 8) return bits * 2.3283064365386963e-10 u = (i / n + offset[0]) % 1.0 v = (radical_inverse(i) + offset[1]) % 1.0 yaw = 2 * math.pi * u pitch = math.asin(2 * v - 1) return yaw, pitch def stable_id(dataset: str, source_model: str) -> str: return f"{dataset}__{source_model}" def load_sources(bench_data_root: Path): metadata_path = bench_data_root / "metadata.json" with metadata_path.open("r", encoding="utf-8") as f: metadata = json.load(f) rows = [] seen = set() for item in metadata: dataset = item["dataset"] source_model = item["source_model"] sid = stable_id(dataset, source_model) if sid in seen: continue seen.add(sid) model_path = bench_data_root / dataset / source_model / "source_model" / "model.glb" rows.append( { "sha256": sid, "dataset": dataset, "source_model": source_model, "local_path": str(model_path), "aesthetic_score": 10.0, "rendered": False, "voxelized": False, } ) return rows def write_metadata(output_dir: Path, rows): output_dir.mkdir(parents=True, exist_ok=True) df = pd.DataFrame(rows) df.to_csv(output_dir / "metadata.csv", index=False) id_map = { row["sha256"]: { "dataset": row["dataset"], "source_model": row["source_model"], "local_path": row["local_path"], } for row in rows } (output_dir / "id_map.json").write_text(json.dumps(id_map, indent=2), encoding="utf-8") def build_views(num_views: int): offset = (0.0, 0.0) # Deterministic preprocessing views. views = [] for i in range(num_views): yaw, pitch = sphere_hammersley_sequence(i, num_views, offset) views.append({"yaw": yaw, "pitch": pitch, "radius": 2, "fov": 40 / 180 * math.pi}) return views def render_one(row, args, views): output_folder = Path(args.output_dir) / "renders" / row["sha256"] transforms = output_folder / "transforms.json" mesh = output_folder / "mesh.ply" if transforms.exists() and mesh.exists() and not args.force_render: return row["sha256"], True, "cached" output_folder.mkdir(parents=True, exist_ok=True) render_script = Path(args.repo_root) / "dataset_toolkits" / "blender_script" / "render.py" cmd = [ args.blender_path, "-b", "-P", str(render_script), "--", "--views", json.dumps(views), "--object", row["local_path"], "--resolution", str(args.resolution), "--output_folder", str(output_folder), "--engine", args.engine, "--save_mesh", ] if args.debug: subprocess.run(cmd, check=True, timeout=args.render_timeout_seconds) else: subprocess.run( cmd, check=True, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, timeout=args.render_timeout_seconds, ) return row["sha256"], transforms.exists() and mesh.exists(), "rendered" def voxelize_one(row, output_dir: Path, force=False): import open3d as o3d import utils3d sha = row["sha256"] voxel_path = output_dir / "voxels" / f"{sha}.ply" if voxel_path.exists() and not force: return sha, True, "cached" mesh_path = output_dir / "renders" / sha / "mesh.ply" if not mesh_path.exists(): return sha, False, "missing mesh" voxel_path.parent.mkdir(parents=True, exist_ok=True) mesh = o3d.io.read_triangle_mesh(str(mesh_path)) vertices = np.clip(np.asarray(mesh.vertices), -0.5 + 1e-6, 0.5 - 1e-6) mesh.vertices = o3d.utility.Vector3dVector(vertices) voxel_grid = o3d.geometry.VoxelGrid.create_from_triangle_mesh_within_bounds( mesh, voxel_size=1 / 64, min_bound=(-0.5, -0.5, -0.5), max_bound=(0.5, 0.5, 0.5), ) points = np.array([voxel.grid_index for voxel in voxel_grid.get_voxels()]) if len(points) == 0: return sha, False, "empty voxel grid" points = (points + 0.5) / 64 - 0.5 utils3d.io.write_ply(str(voxel_path), points) return sha, voxel_path.exists(), "voxelized" def update_metadata_status(output_dir: Path): metadata_path = output_dir / "metadata.csv" df = pd.read_csv(metadata_path) rendered = [] voxelized = [] num_voxels = [] for row in df.to_dict("records"): render_ok = (output_dir / "renders" / row["sha256"] / "transforms.json").exists() voxel_path = output_dir / "voxels" / f"{row['sha256']}.ply" rendered.append(bool(render_ok)) voxelized.append(bool(voxel_path.exists())) if voxel_path.exists(): try: import utils3d points = utils3d.io.read_ply(str(voxel_path))[0] num_voxels.append(int(len(points))) except Exception: num_voxels.append(0) else: num_voxels.append(0) df["rendered"] = rendered df["voxelized"] = voxelized df["num_voxels"] = num_voxels df.to_csv(metadata_path, index=False) def main(): parser = argparse.ArgumentParser() parser.add_argument("--bench_data_root", required=True) parser.add_argument("--output_dir", required=True) parser.add_argument("--repo_root", default="/mnt/zsn/zsn_workspace/3DEditFormer") parser.add_argument("--blender_path", default="/opt/blender-4.2.19-linux-x64/blender") parser.add_argument("--num_views", type=int, default=150) parser.add_argument("--resolution", type=int, default=512) parser.add_argument("--engine", default="CYCLES") parser.add_argument("--max_workers", type=int, default=1) parser.add_argument("--max_items", type=int, default=None) parser.add_argument("--render_timeout_seconds", type=int, default=1800) parser.add_argument("--skip_render", action="store_true") parser.add_argument("--skip_voxelize", action="store_true") parser.add_argument("--force_render", action="store_true") parser.add_argument("--force_voxelize", action="store_true") parser.add_argument("--debug", action="store_true") args = parser.parse_args() bench_data_root = Path(args.bench_data_root) output_dir = Path(args.output_dir) rows = load_sources(bench_data_root) if args.max_items is not None: rows = rows[: args.max_items] write_metadata(output_dir, rows) print(f"Prepared metadata for {len(rows)} source models at {output_dir / 'metadata.csv'}", flush=True) views = build_views(args.num_views) if not args.skip_render: print(f"Rendering {len(rows)} source models with {args.num_views} views...", flush=True) with ThreadPoolExecutor(max_workers=args.max_workers) as executor: futures = [executor.submit(render_one, row, args, views) for row in rows] for fut in as_completed(futures): sha, ok, status = fut.result() print(f"render {sha}: {status} ok={ok}", flush=True) if not args.skip_voxelize: print("Voxelizing rendered source meshes...", flush=True) for row in rows: sha, ok, status = voxelize_one(row, output_dir, force=args.force_voxelize) print(f"voxelize {sha}: {status} ok={ok}", flush=True) update_metadata_status(output_dir) print("Done. Next: extract_feature.py, encode_ss_latent.py, encode_latent.py", flush=True) if __name__ == "__main__": main()