nksr_wrapper/reconstructor.py

"""
Core NKSR wrapper: high-level mesh reconstruction API.
"""

from __future__ import annotations

from dataclasses import dataclass
from pathlib import Path
from typing import Optional, Union, Callable
import warnings

import numpy as np
import torch


try:
    import nksr
except ImportError as exc:
    raise ImportError(
        "The `nksr` package is required but not installed. "
        "Please install it from https://github.com/nv-tlabs/NKSR:\n"
        "  git clone https://github.com/nv-tlabs/NKSR.git\n"
        "  cd NKSR && pip install --no-build-isolation package/\n"
        "See the README for environment setup details."
    ) from exc


@dataclass
class MeshResult:
    """Result container for a reconstructed mesh."""

    vertices: np.ndarray
    """(V, 3) float array of mesh vertex positions."""

    faces: np.ndarray
    """(F, 3) int array of triangle face indices."""

    vertex_colors: Optional[np.ndarray] = None
    """(V, 3) float array of per-vertex colors, if texture was reconstructed."""

    def save(self, path: Union[str, Path]) -> None:
        """Save the mesh to a file using Trimesh."""
        import trimesh

        mesh = trimesh.Trimesh(
            vertices=self.vertices,
            faces=self.faces,
            vertex_colors=self.vertex_colors,
        )
        mesh.export(str(path))


class NKSRMeshReconstructor:
    """
    High-level wrapper around the NKSR reconstructor.

    This class hides the internal complexity of NKSR and exposes a single
    ``reconstruct()`` call that takes a point cloud (with optional normals)
    and returns a watertight triangle mesh.

    Parameters
    ----------
    device : str or torch.device, optional
        PyTorch device to run inference on. Default ``"cuda:0"``.
    config : str, optional
        NKSR model configuration to load.  Default ``"ks"`` (kitchen-sink,
        general-purpose pretrained model).  Other options include ``"snet"``
        (ShapeNet objects with normals) and ``"snet-wonormal"`` (ShapeNet
        without normals).
    chunk_tmp_device : str or torch.device, optional
        Temporary offload device for finished chunks when reconstructing very
        large scenes.  Default ``"cpu"``.  Set to ``None`` to disable
        off-loading (keeps everything on *device*).
    """

    def __init__(
        self,
        device: Union[str, torch.device] = "cuda:0",
        config: str = "ks",
        chunk_tmp_device: Optional[Union[str, torch.device]] = "cpu",
    ):
        self.device = torch.device(device)
        self.reconstructor = nksr.Reconstructor(self.device, config=config)

        if chunk_tmp_device is not None:
            self.reconstructor.chunk_tmp_device = torch.device(chunk_tmp_device)

        self._config_name = config

    # ------------------------------------------------------------------ #
    #  Public API                                                        #
    # ------------------------------------------------------------------ #

    def reconstruct(
        self,
        points: np.ndarray,
        normals: Optional[np.ndarray] = None,
        sensor_positions: Optional[np.ndarray] = None,
        colors: Optional[np.ndarray] = None,
        *,
        detail_level: float = 1.0,
        voxel_size: Optional[float] = None,
        chunk_size: float = -1.0,
        overlap_ratio: float = 0.05,
        approx_kernel_grad: bool = False,
        solver_max_iter: int = 2000,
        solver_tol: float = 1e-5,
        nystrom_min_depth: int = 100,
        fused_mode: bool = True,
        mise_iter: int = 1,
        estimate_normals_if_missing: bool = True,
        normal_knn: int = 64,
        normal_drop_threshold_deg: float = 85.0,
    ) -> MeshResult:
        """
        Reconstruct a watertight mesh from a point cloud.

        Parameters
        ----------
        points : np.ndarray
            (N, 3) array of point positions.
        normals : np.ndarray, optional
            (N, 3) array of **oriented** point normals.  If ``None`` and
            *sensor_positions* are also ``None``, normals are estimated on
            the fly (requires *estimate_normals_if_missing* = ``True``).
        sensor_positions : np.ndarray, optional
            (N, 3) array of per-point sensor/camera positions.  When normals
            are missing, NKSR can infer orientation from the point-to-sensor
            vector using the internal ``get_estimate_normal_preprocess_fn``.
        colors : np.ndarray, optional
            (N, 3) array of RGB colors in ``[0, 255]`` or ``[0, 1]``.  If
            provided, the returned mesh will contain per-vertex colors.
        detail_level : float, default 1.0
            Trade-off between smoothness and detail.  ``0.0`` = very smooth,
            ``1.0`` = maximum detail (may over-fit noise).  Ignored when
            *chunk_size* > 0 or *voxel_size* is set.
        voxel_size : float, optional
            Explicit voxel size controlling the reconstruction resolution.
            Overrides *detail_level*.
        chunk_size : float, default -1.0
            Spatial extent of each chunk for out-of-core reconstruction.
            ``-1.0`` disables chunking (process everything at once).  Positive
            values are required for very large point clouds (> few million
            points) to avoid out-of-memory errors.
        overlap_ratio : float, default 0.05
            Overlap between adjacent chunks (as a fraction of *chunk_size*).
        approx_kernel_grad : bool, default False
            Whether to approximate kernel gradients — slightly faster but a
            bit less accurate.
        solver_max_iter : int, default 2000
            Maximum iterations for the sparse PCG linear solver.
        solver_tol : float, default 1e-5
            Convergence tolerance for the PCG solver.
        nystrom_min_depth : int, default 100
            Minimum depth for the Nyström low-rank approximation used by the
            kernel field.
        fused_mode : bool, default True
            Memory-efficient fusion mode when chunking is enabled.
        mise_iter : int, default 1
            Number of MISE (Multi-resolution IsoSurface Extraction) iterations.
            ``0`` = base grid resolution, each additional iteration doubles
            the effective resolution in subdivided cells.
        estimate_normals_if_missing : bool, default True
            If ``True`` and no normals are provided, estimate them from the
            local geometry.  This only works well when the surface is
            sufficiently sampled.
        normal_knn : int, default 64
            k-NN neighborhood size for on-the-fly normal estimation.
        normal_drop_threshold_deg : float, default 85.0
            Maximum angle (in degrees) between the estimated normal and the
            point-to-sensor vector.  Points exceeding this are dropped.

        Returns
        -------
        MeshResult
            Container with ``vertices``, ``faces``, and optionally
            ``vertex_colors``.

        Notes
        -----
        1. **Normals matter.**  NKSR is designed for oriented normals.  If
           your input lacks them, the wrapper will try to estimate them, but
           orientation may be arbitrary (leading to inside-out meshes).
           Providing *sensor_positions* gives the best auto-orientation.
        2. **Scale.**  The default ``voxel_size`` in the ``"ks"`` config is
           ``0.1``.  If your point cloud is in millimetres and represents a
           room-scale scene, ``0.1`` = 10 cm, which is reasonable.  Adjust
           *voxel_size* or scale your data accordingly.
        3. **Chunking.**  When ``chunk_size > 0``, *detail_level* and
           *voxel_size* are ignored by the underlying NKSR code.  To control
           detail in chunked mode, pre-scale the point cloud by
           ``0.1 / desired_voxel_size``.
        """
        points = self._to_tensor(points, "points")

        # ---- handle normals ------------------------------------------------
        preprocess_fn: Optional[Callable] = None

        if normals is not None:
            normals = self._to_tensor(normals, "normals")
        elif sensor_positions is not None:
            sensor_positions = self._to_tensor(sensor_positions, "sensor_positions")
            preprocess_fn = nksr.get_estimate_normal_preprocess_fn(
                knn=normal_knn,
                drop_threshold_degrees=normal_drop_threshold_deg,
            )
        elif estimate_normals_if_missing:
            warnings.warn(
                "No normals or sensor positions provided. "
                "Estimating normals from geometry — orientation may be arbitrary. "
                "Consider providing sensor_positions for best results.",
                UserWarning,
            )
            normals = self._estimate_normals_from_points(points, normal_knn)

        # ---- colors ---------------------------------------------------------
        color_tensor: Optional[torch.Tensor] = None
        if colors is not None:
            colors = np.asarray(colors)
            if colors.max() > 1.0:
                colors = colors / 255.0
            color_tensor = self._to_tensor(colors, "colors")

        # ---- reconstruct ----------------------------------------------------
        field = self.reconstructor.reconstruct(
            xyz=points,
            normal=normals,
            sensor=sensor_positions,
            detail_level=detail_level,
            voxel_size=voxel_size,
            chunk_size=chunk_size,
            overlap_ratio=overlap_ratio,
            approx_kernel_grad=approx_kernel_grad,
            solver_max_iter=solver_max_iter,
            solver_tol=solver_tol,
            nystrom_min_depth=nystrom_min_depth,
            fused_mode=fused_mode,
            preprocess_fn=preprocess_fn,
        )

        # ---- optional texture ------------------------------------------------
        if color_tensor is not None:
            field.set_texture_field(nksr.fields.PCNNField(points, color_tensor))
            if mise_iter < 2:
                warnings.warn(
                    "Color reconstruction requested but mise_iter < 2. "
                    "Increasing to 2 for better color resolution.",
                    UserWarning,
                )
                mise_iter = 2

        # ---- extract mesh ---------------------------------------------------
        mesh = field.extract_dual_mesh(mise_iter=mise_iter)

        vertices = mesh.v.cpu().numpy() if hasattr(mesh.v, "cpu") else np.asarray(mesh.v)
        faces = mesh.f.cpu().numpy() if hasattr(mesh.f, "cpu") else np.asarray(mesh.f)

        vertex_colors = None
        if hasattr(mesh, "c") and mesh.c is not None:
            vertex_colors = (
                mesh.c.cpu().numpy() if hasattr(mesh.c, "cpu") else np.asarray(mesh.c)
            )

        return MeshResult(
            vertices=vertices,
            faces=faces,
            vertex_colors=vertex_colors,
        )

    # ------------------------------------------------------------------ #
    #  Helpers                                                           #
    # ------------------------------------------------------------------ #

    def _to_tensor(self, arr: np.ndarray, name: str) -> torch.Tensor:
        """Convert a numpy array to a float tensor on the target device."""
        arr = np.asarray(arr)
        if arr.ndim != 2 or arr.shape[1] != 3:
            raise ValueError(
                f"{name} must have shape (N, 3), got {arr.shape}"
            )
        return torch.from_numpy(arr).float().to(self.device)

    def _estimate_normals_from_points(
        self, points: torch.Tensor, k: int = 64
    ) -> torch.Tensor:
        """
        Fast PCA-based normal estimation using PyTorch (no Open3D dependency).

        This estimates **unoriented** normals.  Orientation is arbitrary,
        so the resulting mesh may be inside-out.
        """
        # Simple k-NN with brute force — acceptable for moderate N (< 100k).
        # For larger clouds the user should pre-compute normals externally.
        N = points.shape[0]
        if N > 100_000:
            warnings.warn(
                f"Point cloud has {N} points; on-the-fly normal estimation "
                f"may be slow. Consider pre-computing normals with Open3D.",
                UserWarning,
            )

        # Build a KD-tree or use brute force — we use a chunked brute-force
        # approach to keep memory reasonable.
        batch_size = 4096
        normals_list = []

        for i in range(0, N, batch_size):
            batch = points[i : i + batch_size]  # (B, 3)
            # pairwise distances to all points
            dists = torch.cdist(batch, points)  # (B, N)
            _, idx = torch.topk(dists, k=min(k, N), dim=-1, largest=False)  # (B, k)
            neighbors = points[idx]  # (B, k, 3)
            centered = neighbors - neighbors.mean(dim=1, keepdim=True)  # (B, k, 3)
            cov = centered.transpose(1, 2) @ centered  # (B, 3, 3)
            # smallest eigenvector = normal
            eigvals, eigvecs = torch.linalg.eigh(cov)
            normal = eigvecs[:, :, 0]  # (B, 3)
            normals_list.append(normal)

        normals = torch.cat(normals_list, dim=0)
        # arbitrary orientation — flip to point roughly outward from centroid
        centroid = points.mean(dim=0, keepdim=True)
        outward = points - centroid
        flip = (normals * outward).sum(dim=-1, keepdim=True) < 0
        normals = torch.where(flip, -normals, normals)
        return normals