| """Prepare paired, canonical rendered images and coarse Gaussian targets. |
| |
| These are mesh-derived proxy targets, NOT pretrained-teacher outputs or photos. |
| Only public, explicitly allowed assets are downloaded. Source meshes stay local. |
| """ |
| import argparse |
| import gzip |
| import hashlib |
| import json |
| import random |
| import shutil |
| import time |
| import urllib.error |
| import urllib.request |
| from pathlib import Path, PurePosixPath |
|
|
| import numpy as np |
| import pyarrow.parquet as pq |
| import trimesh |
| from huggingface_hub import hf_hub_download |
|
|
| from image3d_student.data import cache_sample, splat_targets |
| from image3d_student.splats import Splats, preview |
| from image3d_student.engine import split_rows |
|
|
| SCENARIO = ("scenario-labs/3d-collection", "55b4c002b4c1529fc48da3dff98516c3069bfa3f") |
| OBJAVERSE = ("allenai/objaverse", "556637099bf4fa79ea7b239d0d9c328b8a2e9ac8") |
|
|
|
|
| def digest(path): |
| h = hashlib.sha256() |
| with Path(path).open("rb") as f: |
| for block in iter(lambda: f.read(1024 * 1024), b""): |
| h.update(block) |
| return h.hexdigest() |
|
|
|
|
| class Downloader: |
| def __init__(self, root, max_file_mb, max_total_mb, reserve_gb): |
| self.root = Path(root) |
| self.root.mkdir(parents=True, exist_ok=True) |
| self.max_file = int(max_file_mb * 1024**2) |
| self.max_total = int(max_total_mb * 1024**2) |
| self.reserve = int(reserve_gb * 1024**3) |
| self.total = 0 |
|
|
| def get(self, repo, revision, filename, expected=None): |
| parts = PurePosixPath(filename) |
| if parts.is_absolute() or ".." in parts.parts or not filename.endswith(".glb"): |
| raise ValueError("Only relative GLB paths are allowed") |
| out = self.root / repo.replace("/", "--") / revision / filename |
| if out.exists(): |
| if expected and digest(out) != expected: |
| raise ValueError("Cached source checksum mismatch") |
| if out.stat().st_size > self.max_file: |
| raise ValueError("Cached asset exceeds per-file limit") |
| return out |
| if self.total >= self.max_total: |
| raise RuntimeError("Download budget exhausted") |
| out.parent.mkdir(parents=True, exist_ok=True) |
| url = f"https://huggingface.co/datasets/{repo}/resolve/{revision}/{filename}" |
| for attempt in range(4): |
| tmp = out.with_suffix(".part") |
| try: |
| |
| with urllib.request.urlopen(url, timeout=60) as response, tmp.open("wb") as f: |
| length = int(response.headers.get("Content-Length", 0)) |
| if length > self.max_file: |
| raise ValueError("Asset exceeds per-file limit") |
| count = 0 |
| while block := response.read(1024 * 1024): |
| count += len(block) |
| self.total += len(block) |
| if count > self.max_file: |
| raise ValueError("Asset exceeds per-file limit") |
| if self.total > self.max_total: |
| raise RuntimeError("Download budget exhausted") |
| if shutil.disk_usage(self.root).free < self.reserve + len(block): |
| raise RuntimeError("Free disk space below reserve") |
| f.write(block) |
| if expected and digest(tmp) != expected: |
| raise ValueError("Source checksum mismatch") |
| tmp.replace(out) |
| return out |
| except urllib.error.HTTPError as exc: |
| if exc.code not in (429, 500, 502, 503, 504) or attempt == 3: |
| raise |
| time.sleep(min(2**attempt, 8)) |
| finally: |
| tmp.unlink(missing_ok=True) |
| raise RuntimeError("Download retries exhausted") |
|
|
|
|
| def scenario_candidates(cache, seed): |
| repo, revision = SCENARIO |
| path = hf_hub_download(repo, "data/train-00000-of-00001.parquet", repo_type="dataset", |
| revision=revision, local_dir=Path(cache) / "metadata/scenario") |
| rows = pq.read_table(path, columns=["id", "license", "model", "inputs", "outputs"]).to_pylist() |
| |
| parent = {} |
|
|
| def find(x): |
| parent.setdefault(x, x) |
| if parent[x] != x: |
| parent[x] = find(parent[x]) |
| return parent[x] |
|
|
| candidates = [] |
| for row in rows: |
| if row["license"] != "cc-by-4.0" or "hunyuan" in (row["model"] or "").lower(): |
| continue |
| assets = (row.get("inputs") or {}).get("assets") or [] |
| keys = ["input:" + (a.get("sha256") or a["path"]) for a in assets if a.get("sha256") or a.get("path")] |
| keys += ["output:" + o["sha256"] for o in row["outputs"] if o.get("sha256")] |
| keys += [row["id"]] |
| for key in keys[1:]: |
| parent[find(key)] = find(keys[0]) |
| for output in row["outputs"]: |
| if output.get("format") == "glb": |
| candidates.append({"repo": repo, "revision": revision, "path": output["path"], |
| "sha256": output.get("sha256"), "asset_id": row["id"], "license": "cc-by-4.0", |
| "author": "Scenario", "title": row["id"], "source_model": row["model"], |
| "source_url": f"https://huggingface.co/datasets/{repo}", "group_key": keys[0]}) |
| break |
| for row in candidates: |
| row["object_id"] = "scenario-" + hashlib.sha256(find(row.pop("group_key")).encode()).hexdigest()[:20] |
| candidates.sort(key=lambda r: r["asset_id"]) |
| random.Random(seed).shuffle(candidates) |
| return candidates |
|
|
|
|
| def objaverse_candidates(cache, seed, shards): |
| repo, revision = OBJAVERSE |
| candidates = [] |
| for shard in range(shards): |
| part = f"000-{shard:03d}" |
| path = hf_hub_download(repo, f"metadata/{part}.json.gz", repo_type="dataset", revision=revision, |
| local_dir=Path(cache) / "metadata/objaverse") |
| with gzip.open(path) as f: |
| records = json.load(f) |
| for uid, row in records.items(): |
| |
| if row.get("license") != "cc0" or row.get("isAgeRestricted") or row.get("animationCount", 0): |
| continue |
| candidates.append({"repo": repo, "revision": revision, "path": f"glbs/{part}/{uid}.glb", |
| "asset_id": uid, "object_id": "objaverse-" + uid, "license": "cc0-1.0", |
| "author": row.get("user", {}).get("displayName", "Unknown"), "title": row.get("name", uid), |
| "source_url": row.get("viewerUrl", f"https://sketchfab.com/3d-models/{uid}")}) |
| candidates.sort(key=lambda r: r["asset_id"]) |
| random.Random(seed).shuffle(candidates) |
| return candidates |
|
|
|
|
| def mesh_splats(path, count, grid, seed): |
| scene = trimesh.load_scene(path, process=False) |
| meshes = [] |
| for node in scene.graph.nodes_geometry: |
| matrix, name = scene.graph[node] |
| mesh = scene.geometry[name].copy() |
| if not isinstance(mesh, trimesh.Trimesh) or not len(mesh.faces): |
| continue |
| mesh.apply_transform(matrix) |
| meshes.append(mesh) |
| if not meshes or sum(len(m.faces) for m in meshes) > 500_000: |
| raise ValueError("Missing mesh or more than 500,000 triangles") |
| bounds = np.array([m.bounds for m in meshes]) |
| low, high = bounds[:, 0].min(axis=0), bounds[:, 1].max(axis=0) |
| extent = float((high - low).max()) |
| if not np.isfinite(bounds).all() or extent < 1e-8: |
| raise ValueError("Degenerate mesh") |
| center = (low + high) / 2 |
| factor = .9 / extent |
| areas = np.array([m.area for m in meshes]) |
| if not np.isfinite(areas).all() or areas.sum() <= 0: |
| raise ValueError("Invalid surface area") |
| counts = np.random.default_rng(seed).multinomial(count, areas / areas.sum()) |
| points, colours = [], [] |
| for mesh, n in zip(meshes, counts): |
| if not n: |
| continue |
| xyz, faces = trimesh.sample.sample_surface(mesh, int(n), seed=seed) |
| bary = trimesh.triangles.points_to_barycentric(mesh.triangles[faces], xyz) |
| vis = mesh.visual |
| if vis.kind == "texture" and vis.uv is not None: |
| uv = (vis.uv[mesh.faces[faces]] * bary[:, :, None]).sum(axis=1) |
| sampled = np.asarray(vis.material.to_color(uv)) |
| if sampled.ndim == 0: |
| raise ValueError("Material has no usable texture or base colour") |
| |
| if sampled.ndim == 1: |
| sampled = np.broadcast_to(sampled, (n, len(sampled))) |
| colour = sampled[:, :3] / 255. |
| elif vis.kind == "vertex": |
| colour = (vis.vertex_colors[mesh.faces[faces], :3] * bary[:, :, None]).sum(axis=1) / 255. |
| elif vis.kind == "face": |
| colour = vis.face_colors[faces, :3] / 255. |
| else: |
| colour = np.full((n, 3), .65) |
| points.append((xyz - center) * factor) |
| colours.append(colour) |
| xyz = np.concatenate(points).astype(np.float32) |
| rgb = np.concatenate(colours).astype(np.float32) |
| n = len(xyz) |
| splats = Splats(xyz, rgb, np.full((n, 3), .6 / grid), |
| np.tile([1, 0, 0, 0], (n, 1)), np.full((n, 1), .9)) |
| transform = np.eye(4) |
| transform[:3, :3] *= factor |
| transform[:3, 3] = -center * factor |
| return splats, transform.tolist() |
|
|
|
|
| def save_progress(out, args, rows, provenance, failures): |
| state = {"version": 1, "settings": vars(args), "rows": rows, |
| "sources": provenance, "skipped": failures} |
| temporary = out / "progress.json.tmp" |
| temporary.write_text(json.dumps(state)) |
| temporary.replace(out / "progress.json") |
|
|
|
|
| def check_resume_settings(previous, args): |
| for key in ("grid", "views", "surface_points", "seed", "scenario", "objaverse"): |
| if previous[key] != getattr(args, key): |
| raise ValueError(f"Cannot change {key} while resuming; use a new dataset directory") |
|
|
|
|
| def recover_legacy(out, args): |
| """Recover the old writer's completed files without rendering or downloading. |
| |
| Old runs saved sources and arrays, but lost manifest rows on insufficient |
| candidates. Recreate those rows and verify targets against cached GLBs. |
| """ |
| previous = json.loads((out / "preparation.json").read_text())["settings"] |
| check_resume_settings(previous, args) |
| provenance = json.loads((out / "sources.json").read_text()) |
| failures = json.loads((out / "skipped.json").read_text()) |
| rows = [] |
| for index, item in enumerate(provenance): |
| path = Path(args.cache) / item["repo"].replace("/", "--") / item["revision"] / item["path"] |
| if not path.is_file() or digest(path) != item["sha256"]: |
| raise ValueError(f"Missing or corrupt cached source for {item['asset_id']}") |
| splats, transform = mesh_splats(path, args.surface_points, args.grid, |
| int(item["sha256"][:8], 16) ^ args.seed) |
| stem = hashlib.sha256((item["repo"] + item["asset_id"]).encode()).hexdigest()[:20] |
| for view in range(args.views): |
| sample_id = f"{stem}-v{view}" |
| image, target = f"images/{sample_id}.png", f"targets/{sample_id}.npz" |
| yaw = 2 * np.pi * view / args.views |
| c, s = np.cos(yaw), np.sin(yaw) |
| rotation = np.array([[c, 0, s], [0, 1, 0], [-s, 0, c]], np.float32) |
| rotated = Splats(splats.xyz @ rotation.T, splats.rgb, splats.scale, splats.quat, splats.opacity) |
| expected = splat_targets(rotated, args.grid) |
| with np.load(out / target, allow_pickle=False) as cached: |
| if set(cached.files) != set(expected): |
| raise ValueError(f"Target fields differ for {sample_id}") |
| for key, value in expected.items(): |
| if not np.array_equal(cached[key], value): |
| raise ValueError(f"Recovered target differs for {sample_id}: {key}") |
| from PIL import Image |
| with Image.open(out / image) as check: |
| check.verify() |
| rows.append({"id": sample_id, "image": image, "target": target, |
| "source": "mesh-proxy:" + item["repo"], "grid": args.grid, |
| "teachers": {"appearance": "mesh-proxy:" + item["repo"], "coarse": None, "fine": None}, |
| "teacher_gaussians": len(splats.xyz), "retained_gaussians": int(expected["occ"].sum()), |
| "geometry_transforms": {"coarse": None, "fine": None}, |
| "image_sha256": digest(out / image), "target_sha256": digest(out / target), |
| "metadata": {**item, "normalization_matrix": transform, "yaw_radians": float(yaw), |
| "render": "orthographic Gaussian proxy, camera +z", |
| "target": "isotropic surface-sampled mesh proxy; not a pretrained teacher"}}) |
| if (index + 1) % 10 == 0 or index + 1 == len(provenance): |
| print(json.dumps({"recovered_objects": index + 1, "existing_objects": len(provenance)}), flush=True) |
| save_progress(out, args, rows, provenance, failures) |
| return rows, provenance, failures |
|
|
|
|
| def load_progress(out, args): |
| if not (out / "progress.json").exists(): |
| return recover_legacy(out, args) |
| state = json.loads((out / "progress.json").read_text()) |
| if state.get("version") != 1: |
| raise ValueError("Unsupported preparation progress version") |
| check_resume_settings(state["settings"], args) |
| for row in state["rows"]: |
| for key in ("image", "target"): |
| path = (out / row[key]).resolve() |
| if not path.is_relative_to(out.resolve()) or digest(path) != row[key + "_sha256"]: |
| raise ValueError(f"Prepared file checksum mismatch: {row['id']} {key}") |
| return state["rows"], state["sources"], state["skipped"] |
|
|
|
|
| def prune_invisible(out, rows, provenance, failures): |
| from PIL import Image |
| bad = set() |
| for row in rows: |
| with Image.open(out / row["image"]) as image: |
| if image.convert("RGB").getbbox() is None: |
| bad.add((row["metadata"]["repo"], row["metadata"]["asset_id"])) |
| for repo, asset in sorted(bad): |
| failures.append({"repo": repo, "asset": asset, |
| "error": "Rendered view has no visible pixels", "terminal": True}) |
| if bad: |
| print(json.dumps({"excluded_invisible_objects": len(bad), "existing_files_preserved": True}), flush=True) |
| return ([r for r in rows if (r["metadata"]["repo"], r["metadata"]["asset_id"]) not in bad], |
| [p for p in provenance if (p["repo"], p["asset_id"]) not in bad]) |
|
|
|
|
| def group_duplicates(rows, provenance): |
| """Link entire objects when any views or target files are byte-identical.""" |
| parent = {} |
| def find(x): |
| parent.setdefault(x, x) |
| if parent[x] != x: |
| parent[x] = find(parent[x]) |
| return parent[x] |
| hashes = {} |
| for row in rows: |
| meta = row["metadata"] |
| original = meta.setdefault("source_object_id", meta["object_id"]) |
| find(original) |
| for field in ("image_sha256", "target_sha256"): |
| key = (field, row[field]) |
| if key in hashes: |
| a, b = find(original), find(hashes[key]) |
| parent[max(a, b)] = min(a, b) |
| else: |
| hashes[key] = original |
| for row in rows: |
| row["metadata"]["object_id"] = find(row["metadata"]["source_object_id"]) |
| for item in provenance: |
| original = item.setdefault("source_object_id", item["object_id"]) |
| item["object_id"] = find(original) |
|
|
|
|
| def prepare(args): |
| out = Path(args.out) |
| occupied = out.exists() and any(out.iterdir()) |
| if occupied and not getattr(args, "resume", False): |
| raise ValueError("Output is not empty. Use --resume to retain completed objects.") |
| if args.scenario + args.objaverse < 4: |
| raise ValueError("Request at least four objects") |
| out.mkdir(parents=True, exist_ok=True) |
| if occupied: |
| all_rows, provenance, failures = load_progress(out, args) |
| else: |
| all_rows, provenance, failures = [], [], [] |
| save_progress(out, args, all_rows, provenance, failures) |
| all_rows, provenance = prune_invisible(out, all_rows, provenance, failures) |
| save_progress(out, args, all_rows, provenance, failures) |
| seen = {item["sha256"] for item in provenance} |
| completed = {(item["repo"], item["asset_id"]) for item in provenance} |
| rejected = {(f["repo"], f["asset"]) for f in failures if f.get("terminal")} |
| counts = {} |
| fetch = Downloader(args.cache, args.max_file_mb, args.max_download_mb, args.reserve_gb) |
| sources = [] |
| if args.scenario: |
| ready = sum(item["repo"] == SCENARIO[0] for item in provenance) |
| candidates = scenario_candidates(args.cache, args.seed) if ready < args.scenario else [] |
| sources.append((SCENARIO[0], candidates, args.scenario, ready)) |
| if args.objaverse: |
| ready = sum(item["repo"] == OBJAVERSE[0] for item in provenance) |
| candidates = objaverse_candidates(args.cache, args.seed, args.objaverse_shards) if ready < args.objaverse else [] |
| sources.append((OBJAVERSE[0], candidates, args.objaverse, ready)) |
| for repo, candidates, wanted, accepted in sources: |
| for item in candidates: |
| if accepted >= wanted: |
| break |
| if (item["repo"], item["asset_id"]) in completed | rejected: |
| continue |
| staging = out / "_staging" |
| try: |
| path = fetch.get(item["repo"], item["revision"], item["path"], item.get("sha256")) |
| checksum = digest(path) |
| if checksum in seen: |
| continue |
| item["sha256"] = checksum |
| item_seed = int(checksum[:8], 16) ^ args.seed |
| splats, transform = mesh_splats(path, args.surface_points, args.grid, item_seed) |
| stem = hashlib.sha256((item["repo"] + item["asset_id"]).encode()).hexdigest()[:20] |
| staging.mkdir(parents=True, exist_ok=True) |
| rows = [] |
| for view in range(args.views): |
| yaw = 2 * np.pi * view / args.views |
| c, s = np.cos(yaw), np.sin(yaw) |
| rotation = np.array([[c, 0, s], [0, 1, 0], [-s, 0, c]], np.float32) |
| rotated = Splats(splats.xyz @ rotation.T, splats.rgb, splats.scale, splats.quat, splats.opacity) |
| image = staging / "input.png" |
| rendered = preview(rotated, size=128) |
| if rendered.getbbox() is None: |
| raise ValueError("Rendered view has no visible pixels") |
| rendered.save(image) |
| metadata = {**item, "object_id": item["object_id"], "normalization_matrix": transform, |
| "yaw_radians": float(yaw), "render": "orthographic Gaussian proxy, camera +z", |
| "target": "isotropic surface-sampled mesh proxy; not a pretrained teacher"} |
| rows.append(cache_sample(staging, f"{stem}-v{view}", image, rotated, args.grid, |
| source="mesh-proxy:" + item["repo"], extra=metadata)) |
| for folder in ("images", "targets"): |
| (out / folder).mkdir(exist_ok=True) |
| for file in (staging / folder).iterdir(): |
| file.replace(out / folder / file.name) |
| all_rows.extend(rows) |
| seen.add(checksum) |
| provenance.append(item) |
| completed.add((item["repo"], item["asset_id"])) |
| accepted += 1 |
| save_progress(out, args, all_rows, provenance, failures) |
| print(json.dumps({"repo": item["repo"], "accepted": accepted, "requested": wanted, |
| "asset": item["asset_id"]}), flush=True) |
| except RuntimeError: |
| raise |
| except Exception as exc: |
| failures.append({"repo": item["repo"], "asset": item["asset_id"], |
| "error": type(exc).__name__ + ": " + str(exc)[:200], |
| "terminal": "no visible pixels" in str(exc)}) |
| print(json.dumps({"skipped": item["asset_id"], "reason": failures[-1]["error"]}), flush=True) |
| finally: |
| shutil.rmtree(staging, ignore_errors=True) |
| counts[repo] = {"requested": wanted, "prepared": accepted} |
| group_duplicates(all_rows, provenance) |
| save_progress(out, args, all_rows, provenance, failures) |
| (out / "sources.json").write_text(json.dumps(provenance, indent=2)) |
| (out / "skipped.json").write_text(json.dumps(failures, indent=2)) |
| report = {"counts": counts, "samples": len(all_rows), "objects": len(provenance), |
| "distinct_groups": len({r["metadata"]["object_id"] for r in all_rows}), |
| "settings": vars(args), "rendered_proxy_data": True, |
| "excluded": ["scenario-labs/3d-collection-hunyuan", "3dlg-hcvc/MultiScan", |
| "huanngzh/3D-Front", "facebook/show3d-dataset", "Voxel51/gaussian_splatting"]} |
| (out / "preparation.json").write_text(json.dumps(report, indent=2)) |
| if any(v["prepared"] < v["requested"] for v in counts.values()): |
| raise ValueError("Not enough eligible assets. Completed objects are saved. Rerun with --resume and a larger --objaverse-shards value.") |
| if report["distinct_groups"] < 4: |
| raise ValueError("Not enough independent object groups") |
| |
| development, test = split_rows(all_rows, .15, args.seed + 1) |
| for name, rows in [("manifest.jsonl", development), ("test-manifest.jsonl", test)]: |
| (out / name).write_text("".join(json.dumps(r) + "\n" for r in rows)) |
| print(json.dumps(report, indent=2)) |
|
|
|
|
| def main(): |
| p = argparse.ArgumentParser(description=__doc__) |
| p.add_argument("--out", required=True) |
| p.add_argument("--cache", default="data/downloads") |
| p.add_argument("--scenario", type=int, default=64) |
| p.add_argument("--objaverse", type=int, default=32) |
| p.add_argument("--objaverse-shards", type=int, default=8) |
| p.add_argument("--resume", action="store_true", help="Reuse and verify completed objects, including older failed runs") |
| p.add_argument("--grid", type=int, choices=[8, 16, 32], default=16) |
| p.add_argument("--views", type=int, default=4) |
| p.add_argument("--surface-points", type=int, default=4096) |
| p.add_argument("--seed", type=int, default=42) |
| p.add_argument("--max-file-mb", type=int, default=16) |
| p.add_argument("--max-download-mb", type=int, default=2048) |
| p.add_argument("--reserve-gb", type=float, default=2) |
| args = p.parse_args() |
| if min(args.scenario, args.objaverse) < 0 or not 1 <= args.objaverse_shards <= 160: |
| p.error("Invalid object counts or shard count") |
| if min(args.views, args.surface_points, args.max_file_mb, args.max_download_mb, args.reserve_gb) <= 0: |
| p.error("Counts and limits must be positive") |
| prepare(args) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
