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