""" models/reconstruction.py ──────────────────────── Three reconstruction loaders, all sharing the same input contract: Inputs (run kwargs) ────────────────────────────────────────────────────────────── image : PIL.Image (RGB source) depth_normalised : np.ndarray (H×W f32) (normalised [0,1]) depth_raw : np.ndarray (H×W f32) (raw values) output_dir : str (directory to write PLY files into) alpha_mask : np.ndarray (H×W f32, [0,1]) optional — foreground matte from the background-removal stage. When present, this replaces the brightness-based sky/background heuristic in rgbd_to_pointcloud(). Outputs (returned dict) ────────────────────────────────────────────────────────────── ply_path : str (path to written .ply) point_count : int model : str stage : "pointcloud" | "gaussian_scaffold" """ from __future__ import annotations import logging import os from pathlib import Path from typing import Any import numpy as np from PIL import Image from models.base_loader import BaseLoader from utils.image_utils import to_numpy, normalise_depth, rgbd_to_pointcloud, save_ply logger = logging.getLogger(__name__) # ── 1. Open3D RGBD → coloured point cloud ───────────────────────────────────── class Open3DReconstructionLoader(BaseLoader): """ Pure geometry reconstruction — no ML weights. Uses camera back-projection via estimated intrinsics. """ def load(self) -> None: # No ML weights — pure numpy/plyfile back-projection, no open3d needed self._loaded = True logger.info("Open3D-style loader ready (pure numpy back-projection)") def run(self, **inputs: Any) -> dict[str, Any]: if not self._loaded: self.load() image: Image.Image = inputs["image"] depth_norm: np.ndarray = inputs["depth_normalised"] output_dir: str = inputs.get("output_dir", "outputs") alpha_mask: np.ndarray | None = inputs.get("alpha_mask") depth_scale: float = float(self.kwargs.get("depth_scale", 1000.0)) rgb = to_numpy(image.convert("RGB")) # Resize depth to match rgb if needed if depth_norm.shape[:2] != rgb.shape[:2]: from PIL import Image as PILImg d_pil = PILImg.fromarray(depth_norm, mode="F") # mode F = float32 d_pil = d_pil.resize((rgb.shape[1], rgb.shape[0]), PILImg.BILINEAR) depth_norm = np.array(d_pil, dtype=np.float32) h, w = rgb.shape[:2] # Input images have no EXIF, so we use a 55° hFOV assumption — # typical for photographic primes. 60° (old value) was too wide and # sheared the geometry. fx = fy = w / (2 * np.tan(np.radians(27.5))) points, colors = rgbd_to_pointcloud( rgb, depth_norm, fx=fx, fy=fy, depth_scale=depth_scale, max_depth=float(depth_scale), fg_mask=alpha_mask, ) ply_path = Path(output_dir) / "pointcloud.ply" save_ply(ply_path, points, colors) logger.info("Point cloud: %d points → %s", len(points), ply_path) return { "ply_path": str(ply_path), "point_count": len(points), "model": "Open3D RGBD back-projection", "stage": "pointcloud", } # ── 2. Gaussian Splat scaffold (CPU initialiser) ────────────────────────────── class GaussianSplatLoader(BaseLoader): """ Converts a point cloud into a 3DGS initialisation PLY. Each Gaussian is seeded with: • position (xyz from point cloud) • colour (SH DC term from RGB) • opacity (init_opacity) • scale (isotropic, estimated from neighbour distances) • rotation (identity quaternion) The output is a valid 3DGS initialisation file compatible with gaussian-splatting training codebases (e.g. graphdeco-inria/gaussian-splatting). """ def load(self) -> None: self._loaded = True logger.info("Gaussian scaffold loader ready (no weights)") def run(self, **inputs: Any) -> dict[str, Any]: if not self._loaded: self.load() image: Image.Image = inputs["image"] depth_norm: np.ndarray = inputs["depth_normalised"] output_dir: str = inputs.get("output_dir", "outputs") alpha_mask: np.ndarray | None = inputs.get("alpha_mask") sh_degree: int = int(self.kwargs.get("sh_degree", 3)) init_opacity: float = float(self.kwargs.get("init_opacity", 0.1)) rgb = to_numpy(image.convert("RGB")) # --- Step 1: back-project to point cloud --- if depth_norm.shape[:2] != rgb.shape[:2]: from PIL import Image as PILImg d_pil = PILImg.fromarray(depth_norm, mode="F") # mode F = float32 d_pil = d_pil.resize((rgb.shape[1], rgb.shape[0]), PILImg.BILINEAR) depth_norm = np.array(d_pil, dtype=np.float32) h, w = rgb.shape[:2] fx = fy = w / (2 * np.tan(np.radians(27.5))) # 55° hFOV points, colors = rgbd_to_pointcloud( rgb, depth_norm, fx=fx, fy=fy, depth_scale=5.0, fg_mask=alpha_mask, ) n = len(points) logger.info("Scaffolding %d Gaussians", n) # --- Step 2: estimate isotropic scale from kNN distances --- scales = _estimate_scales(points, k=6) # --- Step 3: build Gaussian attribute arrays --- # SH DC coefficient (RGB → SH DC via C0 = 0.28209) SH_C0 = 0.28209479177387814 rgb_float = colors.astype(np.float32) / 255.0 sh_dc = (rgb_float - 0.5) / SH_C0 # inverse of SH→RGB # Opacity in logit space (sigmoid inverse) raw_opacity = np.log(init_opacity / (1 - init_opacity)) * np.ones(n, dtype=np.float32) # Identity rotation as quaternion (w, x, y, z) rotations = np.tile([1.0, 0.0, 0.0, 0.0], (n, 1)).astype(np.float32) # --- Step 4: write 3DGS PLY --- ply_path = Path(output_dir) / "gaussian_scaffold.ply" _write_gaussian_ply(ply_path, points, sh_dc, raw_opacity, scales, rotations) return { "ply_path": str(ply_path), "point_count": n, "model": "Gaussian scaffold initialiser", "stage": "gaussian_scaffold", } # ── 3. DepthSplat (feed-forward, requires GPU) ─────────────────────────────── class DepthSplatLoader(BaseLoader): """ Loads haofeixu/depthsplat from HuggingFace and runs feed-forward Gaussian prediction. Requires CUDA. NOTE: This is a scaffold — the actual DepthSplat model API may need adjustment to match the official release. Pin the model version in requirements.txt once the HF Space is stable. """ def load(self) -> None: if self._loaded: return if self.device.type == "cpu": raise RuntimeError("DepthSplatLoader requires CUDA. Switch to Open3D or Gaussian scaffold on CPU.") logger.info("Loading DepthSplat from HF: %s", self.model_id) # NOTE: DepthSplat does not yet have a diffusers-style HF pipeline. # When the official weights/API land, replace the stub below. try: from huggingface_hub import snapshot_download self._weights_dir = snapshot_download(self.model_id) logger.info("DepthSplat weights cached at %s", self._weights_dir) except Exception as e: raise RuntimeError(f"Could not download DepthSplat weights: {e}") from e self._loaded = True def run(self, **inputs: Any) -> dict[str, Any]: if not self._loaded: self.load() # Placeholder: route back to Gaussian scaffold until official API logger.warning("DepthSplat run() is a scaffold — routing to GaussianSplatLoader") fallback = GaussianSplatLoader( model_id="__builtin_gaussian__", device=self.device, sh_degree=3, init_opacity=0.1, ) fallback.load() return fallback.run(**inputs) # ── Private helpers ─────────────────────────────────────────────────────────── def _estimate_scales(points: np.ndarray, k: int = 6) -> np.ndarray: """ Estimate per-Gaussian isotropic scale as mean distance to k nearest neighbours. Two hard caps prevent a handful of outlier points (which have enormous neighbour distances) from producing Gaussians that smear across the whole scene: • lower cap : 1e-4 (avoids degenerate zero-scale Gaussians) • upper cap : 1% of the scene's bounding-box diagonal — a Gaussian larger than this is almost certainly noise, not geometry. Falls back to a global median if scipy is unavailable. """ # Scene-fraction upper cap bbox_diag = float(np.linalg.norm(points.max(axis=0) - points.min(axis=0))) scale_max = max(bbox_diag * 0.01, 1e-3) try: from scipy.spatial import KDTree tree = KDTree(points) dists, _ = tree.query(points, k=k + 1) # first result is self (dist=0) mean_dist = dists[:, 1:].mean(axis=1).astype(np.float32) clamped = np.clip(mean_dist, 1e-4, scale_max) return clamped.reshape(-1, 1).repeat(3, axis=1) except Exception: global_scale = float(np.median(np.linalg.norm(points, axis=1)) * 0.005) return np.full((len(points), 3), max(global_scale, 1e-4), dtype=np.float32) def _write_gaussian_ply( path: Path, xyz: np.ndarray, # N×3 sh_dc: np.ndarray, # N×3 (DC SH coefficients) opacity: np.ndarray, # N scales: np.ndarray, # N×3 rotations: np.ndarray, # N×4 (w, x, y, z) ) -> None: """Write a 3DGS-compatible PLY with the standard attribute layout.""" from plyfile import PlyData, PlyElement path.parent.mkdir(parents=True, exist_ok=True) n = len(xyz) dtype_fields = ( [("x", "f4"), ("y", "f4"), ("z", "f4")] + [("nx", "f4"), ("ny", "f4"), ("nz", "f4")] # normals (zero) + [(f"f_dc_{i}", "f4") for i in range(3)] + [("opacity", "f4")] + [(f"scale_{i}", "f4") for i in range(3)] + [(f"rot_{i}", "f4") for i in range(4)] ) vertex = np.zeros(n, dtype=dtype_fields) vertex["x"], vertex["y"], vertex["z"] = xyz[:, 0], xyz[:, 1], xyz[:, 2] vertex["nx"] = vertex["ny"] = vertex["nz"] = 0.0 for i in range(3): vertex[f"f_dc_{i}"] = sh_dc[:, i] vertex["opacity"] = opacity for i in range(3): vertex[f"scale_{i}"] = np.log(scales[:, i]) # 3DGS stores log-scale for i in range(4): vertex[f"rot_{i}"] = rotations[:, i] el = PlyElement.describe(vertex, "vertex") PlyData([el], byte_order="<").write(str(path)) logger.info("Gaussian scaffold PLY written → %s (%d Gaussians)", path, n)