| |
| |
| import hashlib |
| import json |
| from pathlib import Path |
|
|
| import numpy as np |
| import torch |
| from PIL import Image, ImageOps |
| from torch.utils.data import Dataset |
|
|
| from .splats import Splats, read_ply |
|
|
|
|
| def grid_centers(grid): |
| axis = (np.arange(grid, dtype=np.float32) + .5) / grid - .5 |
| return np.stack(np.meshgrid(axis, axis, axis, indexing="ij"), axis=-1).reshape(-1, 3) |
|
|
|
|
| def cell_indices(xyz, grid): |
| if not np.isfinite(xyz).all() or (np.abs(xyz) > .50001).any(): |
| raise ValueError("Targets must be aligned inside [-0.5,0.5]^3. No silent normalization is performed.") |
| ijk = np.floor((xyz + .5) * grid).astype(np.int64).clip(0, grid-1) |
| return (ijk[:, 0] * grid + ijk[:, 1]) * grid + ijk[:, 2] |
|
|
|
|
| def splat_targets(splats, grid): |
| """One highest-opacity teacher Gaussian per occupied cell; explicitly lossy.""" |
| cells = cell_indices(splats.xyz, grid) |
| order = np.argsort(-splats.opacity[:, 0], kind="stable") |
| _, first = np.unique(cells[order], return_index=True) |
| keep = order[first] |
| slots = cells[keep] |
| result = {"occ": np.zeros((grid**3, 1), np.float32), "xyz": grid_centers(grid), |
| "rgb": np.zeros((grid**3, 3), np.float32), |
| "scale": np.full((grid**3, 3), 1 / grid / 3, np.float32), |
| "quat": np.tile(np.array([1, 0, 0, 0], np.float32), (grid**3, 1)), |
| "opacity": np.zeros((grid**3, 1), np.float32), |
| "geometry": np.zeros((grid**3, 2), np.float32), |
| "geometry_valid": np.zeros(2, np.float32), "grid": np.array(grid, np.int32)} |
| result["occ"][slots] = 1 |
| for key in ("xyz", "rgb", "scale", "quat", "opacity"): |
| result[key][slots] = getattr(splats, key)[keep] |
| return result |
|
|
|
|
| def mesh_occupancy(path, transform_path, grid): |
| import trimesh |
| if not transform_path: |
| raise ValueError("Every geometry teacher needs an explicit --*-transform JSON matrix, even for identity.") |
| transform = np.asarray(json.loads(Path(transform_path).read_text()), dtype=float) |
| if transform.shape != (4, 4) or not np.isfinite(transform).all() or not np.allclose(transform[3], [0,0,0,1]): |
| raise ValueError("Transform must be a finite affine 4x4 matrix") |
| mesh = trimesh.load(str(path), force="mesh", process=False) |
| if not isinstance(mesh, trimesh.Trimesh) or len(mesh.faces) == 0: |
| raise ValueError("Geometry teacher must contain a triangle mesh") |
| mesh.apply_transform(transform) |
| cell_indices(np.asarray(mesh.vertices), grid) |
| points, _ = trimesh.sample.sample_surface(mesh, min(200_000, grid**3 * 16), seed=42) |
| mask = np.zeros(grid**3, np.float32) |
| mask[np.unique(cell_indices(points, grid))] = 1 |
| return mask |
|
|
|
|
| def cache_sample(root, sample_id, image_path, splats, grid=16, source="idol", |
| coarse=None, coarse_transform=None, fine=None, fine_transform=None, extra=None): |
| root = Path(root) |
| if not sample_id or any(c not in "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789-_" for c in sample_id): |
| raise ValueError("Sample ID may contain only letters, digits, - and _") |
| if grid not in (8, 16, 32): |
| raise ValueError("Grid must be 8, 16 or 32") |
| target = splat_targets(splats, grid) |
| for channel, mesh, transform in [(0, coarse, coarse_transform), (1, fine, fine_transform)]: |
| if mesh: |
| target["geometry"][:, channel] = mesh_occupancy(mesh, transform, grid) |
| target["geometry_valid"][channel] = 1 |
| row = {"id": sample_id, "image": f"images/{sample_id}.png", "target": f"targets/{sample_id}.npz", |
| "source": source, "grid": grid, "teachers": {"appearance": source, |
| "coarse": "InstantMesh" if coarse else None, "fine": "Hi3D" if fine else None}, |
| "teacher_gaussians": len(splats.xyz), "retained_gaussians": int(target["occ"].sum()), |
| "geometry_transforms": {"coarse": json.loads(Path(coarse_transform).read_text()) if coarse else None, |
| "fine": json.loads(Path(fine_transform).read_text()) if fine else None}, |
| "metadata": extra or {}} |
| manifest = root / "manifest.jsonl" |
| rows = load_manifest(manifest) if manifest.exists() else [] |
| if any(r["id"] == sample_id for r in rows): |
| raise ValueError(f"Sample {sample_id} already exists; choose a new ID") |
| if rows and any(r["grid"] != grid for r in rows): |
| raise ValueError("Do not mix grid resolutions in one dataset") |
| image = Image.open(image_path) |
| image.load() |
| if root.joinpath(row["image"]).exists() or root.joinpath(row["target"]).exists(): |
| raise ValueError("Sample files already exist; refusing to overwrite") |
| root.joinpath("images").mkdir(parents=True, exist_ok=True) |
| root.joinpath("targets").mkdir(parents=True, exist_ok=True) |
| image.save(root / row["image"]) |
| row["image_sha256"] = hashlib.sha256((root / row["image"]).read_bytes()).hexdigest() |
| np.savez_compressed(root / row["target"], **target) |
| row["target_sha256"] = hashlib.sha256((root / row["target"]).read_bytes()).hexdigest() |
| |
| tmp = manifest.with_suffix(".tmp") |
| tmp.write_text("".join(json.dumps(r) + "\n" for r in [*rows, row])) |
| tmp.replace(manifest) |
| return row |
|
|
|
|
| def load_manifest(path): |
| rows = [json.loads(line) for line in Path(path).read_text().splitlines() if line.strip()] |
| if not rows or len({r["id"] for r in rows}) != len(rows): |
| raise ValueError("Manifest must contain nonempty, unique sample IDs") |
| return rows |
|
|
|
|
| def image_tensor(path, size=128): |
| """Already-framed RGB or RGBA cutout; black background, square letterbox.""" |
| with Image.open(path) as original: |
| rgba = ImageOps.exif_transpose(original).convert("RGBA") |
| bg = Image.new("RGBA", rgba.size, (0, 0, 0, 255)) |
| bg.alpha_composite(rgba) |
| image = ImageOps.pad(bg.convert("RGB"), (size, size), method=Image.Resampling.LANCZOS, color="black") |
| return torch.from_numpy(np.array(image, dtype=np.float32).transpose(2, 0, 1) / 255) |
|
|
|
|
| class CachedDataset(Dataset): |
| def __init__(self, manifest, rows=None): |
| self.root = Path(manifest).resolve().parent |
| self.rows = load_manifest(manifest) if rows is None else rows |
| if not self.rows: |
| raise ValueError("Empty dataset") |
| self.grid = self.rows[0]["grid"] |
| for row in self.rows: |
| if row["grid"] != self.grid: |
| raise ValueError("Mixed grids") |
| for key in ("image", "target"): |
| p = (self.root / row[key]).resolve() |
| if not p.is_relative_to(self.root) or not p.is_file(): |
| raise ValueError(f"Invalid dataset {key} path") |
| digest = hashlib.sha256(p.read_bytes()).hexdigest() |
| if digest != row[f"{key}_sha256"]: |
| raise ValueError(f"Dataset {key} hash mismatch for {row['id']}") |
|
|
| def __len__(self): |
| return len(self.rows) |
|
|
| def __getitem__(self, i): |
| row = self.rows[i] |
| with np.load(self.root / row["target"], allow_pickle=False) as cached: |
| target = {k: torch.from_numpy(cached[k].copy()).float() for k in cached.files if k != "grid"} |
| return image_tensor(self.root / row["image"]), target |
|
|
