Upload nksr_wrapper/reconstructor.py

nksr_wrapper/reconstructor.py

@@ -0,0 +1,326 @@
+"""
+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