models/reconstruction.py

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