diff --git a/surface_reconstruction/nksr/.gitignore b/surface_reconstruction/nksr/.gitignore new file mode 100644 index 0000000..e78c6f4 --- /dev/null +++ b/surface_reconstruction/nksr/.gitignore @@ -0,0 +1,177 @@ +# Byte-compiled / optimized / DLL files +__pycache__/ +*.py[cod] +*$py.class + +# C extensions +*.so + +# Distribution / packaging +.Python +build/ +develop-eggs/ +dist/ +downloads/ +eggs/ +.eggs/ +lib/ +lib64/ +parts/ +sdist/ +var/ +wheels/ +share/python-wheels/ +*.egg-info/ +.installed.cfg +*.egg +MANIFEST + +# PyInstaller +# Usually these files are written by a python script from a template +# before PyInstaller builds the exe, so as to inject date/other infos into it. +*.manifest +*.spec + +# Installer logs +pip-log.txt +pip-delete-this-directory.txt + +# Unit test / coverage reports +htmlcov/ +.tox/ +.nox/ +.coverage +.coverage.* +.cache +nosetests.xml +coverage.xml +*.cover +*.py,cover +.hypothesis/ +.pytest_cache/ +cover/ + +# Translations +*.mo +*.pot + +# Django stuff: +*.log +local_settings.py +db.sqlite3 +db.sqlite3-journal + +# Flask stuff: +instance/ +.webassets-cache + +# Scrapy stuff: +.scrapy + +# Sphinx documentation +docs/_build/ + +# PyBuilder +.pybuilder/ +target/ + +# Jupyter Notebook +.ipynb_checkpoints + +# IPython +profile_default/ +ipython_config.py + +# pyenv +# For a library or package, you might want to ignore these files since the code is +# intended to run in multiple environments; otherwise, check them in: +# .python-version + +# pipenv +# According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. +# However, in case of collaboration, if having platform-specific dependencies or dependencies +# having no cross-platform support, pipenv may install dependencies that don't work, or not +# install all needed dependencies. +#Pipfile.lock + +# poetry +# Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control. +# This is especially recommended for binary packages to ensure reproducibility, and is more +# commonly ignored for libraries. +# https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control +#poetry.lock + +# PEP 582; used by e.g. github.com/David-OConnor/pyflow +__pypackages__/ + +# Celery stuff +celerybeat-schedule +celerybeat.pid + +# SageMath parsed files +*.sage.py + +# Environments +.env +.venv +env/ +venv/ +ENV/ +env.bak/ +venv.bak/ + +# Spyder project settings +.spyderproject +.spyproject + +# Rope project settings +.ropeproject + +# mkdocs documentation +/site + +# mypy +.mypy_cache/ +.dmypy.json +dmypy.json + +# Pyre type checker +.pyre/ + +# pytype static type analyzer +.pytype/ + +# Cython debug symbols +cython_debug/ + +# PyCharm +# JetBrains specific template is maintainted in a separate JetBrains.gitignore that can +# be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore +# and can be added to the global gitignore or merged into this file. For a more nuclear +# option (not recommended) you can uncomment the following to ignore the entire idea folder. +.idea/ + +cmake-build-debug +cmake-build-release + +.premature_checkpoints + +*.old +sync-pycg.sh +sync-dataloader.sh +ngc_mapping.json +playground/ + +# Sweep-related files +sweeps/ +wandb/ +wdb_sweep.py + +zeus_config.yaml +assets/buda.ply +assets/waymo-pcd.ply +assets/scannet.ply +assets/las/ + +package/external/ +.vscode/ diff --git a/surface_reconstruction/nksr/LICENSE b/surface_reconstruction/nksr/LICENSE new file mode 100644 index 0000000..c73d0d3 --- /dev/null +++ b/surface_reconstruction/nksr/LICENSE @@ -0,0 +1,91 @@ + +Copyright (c) 2023, NVIDIA Corporation & affiliates. All rights reserved. + + +NVIDIA Source Code License for NKSR + + +======================================================================= + +1. Definitions + +“Licensor” means any person or entity that distributes its Work. + +“Work” means (a) the original work of authorship made available under +this license, which may include software, documentation, or other files, +and (b) any additions to or derivative works thereof that are made +available under this license. + +The terms “reproduce,” “reproduction,” “derivative works,” and +“distribution” have the meaning as provided under U.S. copyright law; +provided, however, that for the purposes of this license, derivative works +shall not include works that remain separable from, or merely link +(or bind by name) to the interfaces of, the Work. + +Works are “made available” under this license by including in or with +the Work either (a) a copyright notice referencing the applicability of +this license to the Work, or (b) a copy of this license. + +2. License Grant + + 2.1 Copyright Grant. Subject to the terms and conditions of this license, + each Licensor grants to you a perpetual, worldwide, non-exclusive, + royalty-free, copyright license to use, reproduce, prepare derivative + works of, publicly display, publicly perform, sublicense and distribute + its Work and any resulting derivative works in any form. + +3. Limitations + + 3.1 Redistribution. You may reproduce or distribute the Work only if + (a) you do so under this license, (b) you include a complete copy of + this license with your distribution, and (c) you retain without + modification any copyright, patent, trademark, or attribution notices + that are present in the Work. + + 3.2 Derivative Works. You may specify that additional or different terms + apply to the use, reproduction, and distribution of your derivative + works of the Work (“Your Terms”) only if (a) Your Terms provide that the + use limitation in Section 3.3 applies to your derivative works, and (b) + you identify the specific derivative works that are subject to Your Terms. + Notwithstanding Your Terms, this license (including the redistribution + requirements in Section 3.1) will continue to apply to the Work itself. + + 3.3 Use Limitation. The Work and any derivative works thereof only may be + used or intended for use non-commercially. Notwithstanding the foregoing, + NVIDIA Corporation and its affiliates may use the Work and any derivative + works commercially. As used herein, “non-commercially” means for research + or evaluation purposes only. + + 3.4 Patent Claims. If you bring or threaten to bring a patent claim against + any Licensor (including any claim, cross-claim or counterclaim in a lawsuit) + to enforce any patents that you allege are infringed by any Work, then your + rights under this license from such Licensor (including the grant in + Section 2.1) will terminate immediately. + + 3.5 Trademarks. This license does not grant any rights to use any Licensor’s + or its affiliates’ names, logos, or trademarks, except as necessary to + reproduce the notices described in this license. + + 3.6 Termination. If you violate any term of this license, then your rights + under this license (including the grant in Section 2.1) will terminate + immediately. + +4. Disclaimer of Warranty. + +THE WORK IS PROVIDED “AS IS” WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, +EITHER EXPRESS OR IMPLIED, INCLUDING WARRANTIES OR CONDITIONS OF +MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE, TITLE OR NON-INFRINGEMENT. +YOU BEAR THE RISK OF UNDERTAKING ANY ACTIVITIES UNDER THIS LICENSE. + +5. Limitation of Liability. + +EXCEPT AS PROHIBITED BY APPLICABLE LAW, IN NO EVENT AND UNDER NO LEGAL THEORY, +WHETHER IN TORT (INCLUDING NEGLIGENCE), CONTRACT, OR OTHERWISE SHALL ANY +LICENSOR BE LIABLE TO YOU FOR DAMAGES, INCLUDING ANY DIRECT, INDIRECT, SPECIAL, +INCIDENTAL, OR CONSEQUENTIAL DAMAGES ARISING OUT OF OR RELATED TO THIS LICENSE, +THE USE OR INABILITY TO USE THE WORK (INCLUDING BUT NOT LIMITED TO LOSS OF +GOODWILL, BUSINESS INTERRUPTION, LOST PROFITS OR DATA, COMPUTER FAILURE OR +MALFUNCTION, OR ANY OTHER DAMAGES OR LOSSES), EVEN IF THE LICENSOR HAS BEEN +ADVISED OF THE POSSIBILITY OF SUCH DAMAGES. + +======================================================================= \ No newline at end of file diff --git a/surface_reconstruction/nksr/README.md b/surface_reconstruction/nksr/README.md new file mode 100644 index 0000000..d15d989 --- /dev/null +++ b/surface_reconstruction/nksr/README.md @@ -0,0 +1,20 @@ +# Neural Kernel Surface Reconstruction (NKSR) + +Idea: This is supposed to be a minimal implementation of NN-based surface reconstruction method. +Since the common belief is that NN-based methods are very scalable with more and more data, we should go with the very brute force network method instead of the more sophisticated ones (i.e. kernel solve). + +## Installation + +Prerequisites: Follow fVDB [official website](https://github.com/openvdb/fvdb-core) to install the conda environment of fVDB. + +```bash +# Editable install +pip install --no-build-isolation -e . +``` + +## Usage + +```bash +# Quick test +python examples/recons_simple.py +``` diff --git a/surface_reconstruction/nksr/examples/recons_simple.py b/surface_reconstruction/nksr/examples/recons_simple.py new file mode 100644 index 0000000..b969e4c --- /dev/null +++ b/surface_reconstruction/nksr/examples/recons_simple.py @@ -0,0 +1,36 @@ +# Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property +# and proprietary rights in and to this software, related documentation +# and any modifications thereto. Any use, reproduction, disclosure or +# distribution of this software and related documentation without an express +# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited. + +from pathlib import Path + +import point_cloud_utils as pcu +import requests +import torch + +import nksr_fvdb + +if __name__ == "__main__": + device = torch.device("cuda:0") + + bunny_mesh_path = Path.home() / "data" / "bunny.ply" + if not bunny_mesh_path.exists(): + bunny_mesh_path.parent.mkdir(parents=True, exist_ok=True) + res = requests.get("https://huggingface.co/heiwang1997/nksr-checkpoints/resolve/main/data/buda.ply") + with bunny_mesh_path.open("wb") as f: + f.write(res.content) + + bunny_v, bunny_n = pcu.load_mesh_vn(bunny_mesh_path) + + input_xyz = torch.from_numpy(bunny_v).float().to(device) + input_normal = torch.from_numpy(bunny_n).float().to(device) + + reconstructor = nksr_fvdb.Reconstructor(device, preset="ks-sdf") + field = reconstructor.reconstruct(input_xyz, input_normal, voxel_size=1.0) + mesh = field.extract_primal_mesh(depth=1) + + # Do something with the mesh... diff --git a/surface_reconstruction/nksr/nksr_fvdb/__init__.py b/surface_reconstruction/nksr/nksr_fvdb/__init__.py new file mode 100644 index 0000000..d508bdf --- /dev/null +++ b/surface_reconstruction/nksr/nksr_fvdb/__init__.py @@ -0,0 +1,12 @@ +# Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property +# and proprietary rights in and to this software, related documentation +# and any modifications thereto. Any use, reproduction, disclosure or +# distribution of this software and related documentation without an express +# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited. + +from nksr_fvdb.reconstructor import Reconstructor + +__version__ = "1.0.0" +__version_info__ = (1, 0, 0) diff --git a/surface_reconstruction/nksr/nksr_fvdb/fields/__init__.py b/surface_reconstruction/nksr/nksr_fvdb/fields/__init__.py new file mode 100644 index 0000000..4ab819e --- /dev/null +++ b/surface_reconstruction/nksr/nksr_fvdb/fields/__init__.py @@ -0,0 +1,12 @@ +# Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property +# and proprietary rights in and to this software, related documentation +# and any modifications thereto. Any use, reproduction, disclosure or +# distribution of this software and related documentation without an express +# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited. + + +from .base_field import BaseField, EvaluationResult +from .layer_field import LayerField +from .neural_field import NeuralField diff --git a/surface_reconstruction/nksr/nksr_fvdb/fields/base_field.py b/surface_reconstruction/nksr/nksr_fvdb/fields/base_field.py new file mode 100644 index 0000000..cd9b5ac --- /dev/null +++ b/surface_reconstruction/nksr/nksr_fvdb/fields/base_field.py @@ -0,0 +1,127 @@ +# Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property +# and proprietary rights in and to this software, related documentation +# and any modifications thereto. Any use, reproduction, disclosure or +# distribution of this software and related documentation without an express +# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited. + + +from abc import ABC +from dataclasses import dataclass +from typing import Optional, Union + +import numpy as np +import torch + +from nksr_fvdb.svh import SparseFeatureHierarchy + + +@dataclass +class EvaluationResult: + value: torch.Tensor + gradient: torch.Tensor | None = None + + @classmethod + def zero(cls, grad: bool = False) -> "EvaluationResult": + return EvaluationResult(torch.tensor(0), torch.tensor(0) if grad else None) + + def __add__(self, other): + assert isinstance(other, EvaluationResult) + return EvaluationResult( + self.value + other.value, (self.gradient + other.gradient) if self.gradient is not None else None + ) + + def __sub__(self, other): + assert not isinstance(other, EvaluationResult) + return EvaluationResult(self.value - other, self.gradient) + + +@dataclass(kw_only=True, slots=True) +class MeshingResult: + vertices: torch.Tensor + faces: torch.Tensor + colors: torch.Tensor | None = None + + +class BaseField(ABC): + """ + Base class for the 3D continuous field: + f_bar = f - level_set + """ + + def __init__(self, svh: Optional[SparseFeatureHierarchy]): + self.svh = svh + self.scale = 1.0 + self.mask_field = None + self.texture_field = None + self.set_level_set(0.0) + + def set_level_set(self, level_set: float): + self.level_set = level_set + + def set_scale(self, scale: float): + self.scale = scale + if self.mask_field is not None: + self.mask_field.set_scale(scale) + if self.texture_field is not None: + self.texture_field.set_scale(scale) + + def to_(self, device: Union[torch.device, str]): + if self.svh is not None: + self.svh.to_(device) + if self.mask_field is not None: + self.mask_field.to_(device) + if self.texture_field is not None: + self.texture_field.to_(device) + + @property + def device(self): + return self.svh.device + + def evaluate_f(self, xyz: torch.Tensor, grad: bool = False): + pass + + def evaluate_f_bar(self, xyz: torch.Tensor, max_points: int = -1, verbose: bool = True): + n_chunks = int(np.ceil(xyz.size(0) / max_points)) if max_points > 0 else 1 + xyz_chunks = torch.chunk(xyz, n_chunks) + f_bar_chunks = [] + + if verbose and len(xyz_chunks) > 10: + from tqdm import tqdm + + xyz_chunks = tqdm(xyz_chunks) + + for xyz_chunk in xyz_chunks: + if self.scale != 1.0: + xyz_chunk = xyz_chunk / self.scale + f_chunk = self.evaluate_f(xyz_chunk, grad=False).value + f_bar_chunks.append(f_chunk - self.level_set) + + return torch.cat(f_bar_chunks) + + def extract_primal_mesh(self, depth: int, resolution: int = 2, trim: bool = True, max_points: int = -1): + primal_grid = self.svh.grids[depth] + primal_grid_dense = primal_grid.subdivided_grid( + resolution, torch.ones(primal_grid.num_voxels, dtype=bool, device=self.svh.device) + ) + dual_grid_dense = primal_grid_dense.dual_grid() + + dual_graph = meshing.primal_cube_graph(primal_grid_dense, dual_grid_dense) + dual_corner_pos = dual_grid_dense.grid_to_world(dual_grid_dense.active_grid_coords().float()) + if self.scale != 1.0: + dual_corner_pos = dual_corner_pos * self.scale + dual_corner_value = self.evaluate_f_bar(dual_corner_pos, max_points=max_points) + + primal_v, primal_f = MarchingCubes().apply(dual_graph, dual_corner_pos, dual_corner_value) + + if self.mask_field is not None and trim: + vert_mask = self.mask_field.evaluate_f_bar(primal_v, max_points=max_points) < 0.0 + primal_v, primal_f = utils.apply_vertex_mask(primal_v, primal_f, vert_mask) + + if self.texture_field is not None: + primal_c = self.texture_field.evaluate_f_bar(primal_v, max_points=max_points) + else: + primal_c = None + + return MeshingResult(primal_v, primal_f, primal_c) diff --git a/surface_reconstruction/nksr/nksr_fvdb/fields/layer_field.py b/surface_reconstruction/nksr/nksr_fvdb/fields/layer_field.py new file mode 100644 index 0000000..7efea8e --- /dev/null +++ b/surface_reconstruction/nksr/nksr_fvdb/fields/layer_field.py @@ -0,0 +1,30 @@ +# Copyright (c) 2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property +# and proprietary rights in and to this software, related documentation +# and any modifications thereto. Any use, reproduction, disclosure or +# distribution of this software and related documentation without an express +# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited. + + +import torch + +from nksr_fvdb.fields.base_field import BaseField, EvaluationResult +from nksr_fvdb.svh import SparseFeatureHierarchy + + +class LayerField(BaseField): + """ + Provides 1 for regions outside layer's voxel boundary, -1 for regions inside (including the exact boundary) + """ + + def __init__(self, svh: SparseFeatureHierarchy, inside_depth: int): + super().__init__(svh) + self.inside_depth = inside_depth + + def evaluate_f(self, xyz: torch.Tensor, grad: bool = False): + assert not grad, "Layer step function does not have valid gradient" + in_grid_mask = self.svh.grids[self.inside_depth - 1].points_in_active_voxel(xyz) + f = torch.ones(xyz.size(0), device=xyz.device) + f[in_grid_mask] = -1 + return EvaluationResult(f) diff --git a/surface_reconstruction/nksr/nksr_fvdb/fields/neural_field.py b/surface_reconstruction/nksr/nksr_fvdb/fields/neural_field.py new file mode 100644 index 0000000..8a01a78 --- /dev/null +++ b/surface_reconstruction/nksr/nksr_fvdb/fields/neural_field.py @@ -0,0 +1,59 @@ +# Copyright (c) 2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property +# and proprietary rights in and to this software, related documentation +# and any modifications thereto. Any use, reproduction, disclosure or +# distribution of this software and related documentation without an express +# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited. + + +from typing import Union + +import torch + +from nksr_fvdb.fields.base_field import BaseField, EvaluationResult +from nksr_fvdb.nn.encdec import MultiscalePointDecoder +from nksr_fvdb.svh import SparseFeatureHierarchy + + +class NeuralField(BaseField): + def __init__( + self, svh: SparseFeatureHierarchy, decoder: MultiscalePointDecoder, features: dict, grad_type: str = "numerical" + ): + super().__init__(svh) + self.decoder = decoder + self.features = features + + assert grad_type in ["numerical", "analytical"] + self.grad_type = grad_type + + def to_(self, device: Union[torch.device, str]): + super().to_(device) + self.features = {k: v.to(device) for k, v in self.features.items()} + + def evaluate_f(self, xyz: torch.Tensor, grad: bool = False): + res = self.decoder(xyz, self.svh, self.features) + assert res.size(1) == 1, "Decoder is only allowed to produce 1-dim vectors." + + if grad: + if self.grad_type == "numerical": + interval = 0.01 * self.svh.voxel_size + grad_value: list[torch.Tensor] = [] + for offset in [(interval, 0, 0), (0, interval, 0), (0, 0, interval)]: + offset_tensor = torch.tensor(offset, device=self.device)[None, :] + res_p = self.decoder(xyz + offset_tensor, self.svh, self.features)[:, 0] + res_n = self.decoder(xyz - offset_tensor, self.svh, self.features)[:, 0] + grad_value.append((res_p - res_n) / (2 * interval)) + grad_value = torch.stack(grad_value, dim=1) + else: + xyz_d = torch.clone(xyz) + xyz_d.requires_grad = True + with torch.enable_grad(): + res_d = self.decoder(xyz_d, self.svh, self.features) + grad_value = torch.autograd.grad( + res_d, [xyz_d], grad_outputs=torch.ones_like(res_d), create_graph=self.decoder.training + )[0] + else: + grad_value = None + + return EvaluationResult(res[:, 0], grad_value) diff --git a/surface_reconstruction/nksr/nksr_fvdb/nn/__init__.py b/surface_reconstruction/nksr/nksr_fvdb/nn/__init__.py new file mode 100644 index 0000000..57acde3 --- /dev/null +++ b/surface_reconstruction/nksr/nksr_fvdb/nn/__init__.py @@ -0,0 +1,11 @@ +# Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property +# and proprietary rights in and to this software, related documentation +# and any modifications thereto. Any use, reproduction, disclosure or +# distribution of this software and related documentation without an express +# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited. + + +from .encdec import MultiscalePointDecoder, PointEncoder +from .unet import SparseStructureNet diff --git a/surface_reconstruction/nksr/nksr_fvdb/nn/encdec.py b/surface_reconstruction/nksr/nksr_fvdb/nn/encdec.py new file mode 100644 index 0000000..d496f19 --- /dev/null +++ b/surface_reconstruction/nksr/nksr_fvdb/nn/encdec.py @@ -0,0 +1,170 @@ +# Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property +# and proprietary rights in and to this software, related documentation +# and any modifications thereto. Any use, reproduction, disclosure or +# distribution of this software and related documentation without an express +# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited. + +import fvdb +import torch +import torch.nn as nn +import torch.nn.functional as F +from fvdb import JaggedTensor +from torch_scatter import scatter_max, scatter_mean + +from nksr_fvdb.svh import SparseFeatureHierarchy + + +class ResnetBlockFC(nn.Module): + """Fully connected ResNet Block class.""" + + def __init__(self, size_in: int, size_out: int | None = None): + super().__init__() + # Attributes + if size_out is None: + size_out = size_in + + size_h = min(size_in, size_out) + self.size_in = size_in + self.size_h = size_h + self.size_out = size_out + # Submodules + self.fc_0 = nn.Linear(size_in, size_h) + self.fc_1 = nn.Linear(size_h, size_out) + self.actvn = nn.ReLU() + + if size_in == size_out: + self.shortcut = None + else: + self.shortcut = nn.Linear(size_in, size_out, bias=False) + # Initialization + nn.init.zeros_(self.fc_1.weight) + + def forward(self, x): + net = self.fc_0(self.actvn(x)) + dx = self.fc_1(self.actvn(net)) + + if self.shortcut is not None: + x_s = self.shortcut(x) + else: + x_s = x + + return x_s + dx + + +class PointEncoder(nn.Module): + """PointNet-based encoder network with ResNet blocks for each point. + Number of input points are fixed. + """ + + def __init__(self, dim: int, c_dim: int = 32, hidden_dim: int = 32, n_blocks: int = 3): + super().__init__() + + self.c_dim = c_dim + self.fc_pos = nn.Linear(dim, 2 * hidden_dim) + self.blocks = nn.ModuleList([ResnetBlockFC(2 * hidden_dim, hidden_dim) for _ in range(n_blocks)]) + self.fc_c = nn.Linear(hidden_dim, c_dim) + self.hidden_dim = hidden_dim + + def forward(self, pts_xyz: JaggedTensor, pts_feature: JaggedTensor, svh: SparseFeatureHierarchy, depth: int = 0): + grid = svh.grids[depth] + + # Get voxel idx + pts_xyz = grid.world_to_voxel(pts_xyz) + vid = grid.ijk_to_index(pts_xyz.round().int()).jdata + + # Map coordinates to local voxel + pts_xyz = pts_xyz.jdata + pts_xyz = (pts_xyz + 0.5) % 1 + + pts_mask = vid != -1 + vid, pts_xyz = vid[pts_mask], pts_xyz[pts_mask] + + # Feature extraction + if pts_feature is None: + pts_feature = self.fc_pos(pts_xyz) + else: + pts_feature = pts_feature.jdata + pts_feature = pts_feature[pts_mask] + pts_feature = self.fc_pos(torch.cat([pts_xyz, pts_feature], dim=1)) + pts_feature = self.blocks[0](pts_feature) + for block in self.blocks[1:]: + pooled, _ = scatter_max(pts_feature, vid, dim=0, dim_size=grid.total_voxels) + pooled = pooled[vid] + pts_feature = torch.cat([pts_feature, pooled], dim=1) + pts_feature = block(pts_feature) + + c = self.fc_c(pts_feature) + c = scatter_mean(c, vid, dim=0, dim_size=grid.total_voxels) + return grid.jagged_like(c) + + +class MultiscalePointDecoder(nn.Module): + def __init__( + self, + c_each_dim: int = 16, + multiscale_depths: int = 4, + p_dim: int = 3, + out_dim: int = 1, + hidden_size: int = 32, + n_blocks: int = 2, + aggregation: str = "cat", + out_init: float = None, + coords_depths: list = None, + ): + if aggregation == "cat": + c_dim = c_each_dim * multiscale_depths + elif aggregation == "sum": + c_dim = c_each_dim + else: + raise NotImplementedError + + if coords_depths is None: + coords_depths = list(range(multiscale_depths)) + coords_depths = sorted(coords_depths) + + super().__init__() + self.c_dim = c_dim + self.c_each_dim = c_each_dim + self.n_blocks = n_blocks + self.multiscale_depths = multiscale_depths + self.aggregation = aggregation + self.coords_depths = coords_depths + + self.fc_c = nn.ModuleList([nn.Linear(c_dim, hidden_size) for _ in range(n_blocks)]) + self.fc_p = nn.Linear(p_dim * len(coords_depths), hidden_size) + self.blocks = nn.ModuleList([ResnetBlockFC(hidden_size) for _ in range(n_blocks)]) + self.fc_out = nn.Linear(hidden_size, out_dim) + self.out_dim = out_dim + + # Init parameters + if out_init is not None: + nn.init.zeros_(self.fc_out.weight) + nn.init.constant_(self.fc_out.bias, out_init) + + def forward(self, xyz: JaggedTensor, svh: SparseFeatureHierarchy, multiscale_feat: dict[int, JaggedTensor]): + p_feats = [] + for did in self.coords_depths: + vs = svh.grids[did].voxel_sizes[xyz.jidx.int()] + p = (xyz.jdata % vs) / vs - 0.5 + p_feats.append(p) + p = torch.cat(p_feats, dim=1) + + c_feats = [] + for did in range(self.multiscale_depths): + c = svh.grids[did].sample_trilinear(xyz, multiscale_feat[did]) + c_feats.append(c.jdata) + + if self.aggregation == "cat": + c = torch.cat(c_feats, dim=1) + else: + c = sum(c_feats) + + net = self.fc_p(p) + for i in range(self.n_blocks): + net = net + self.fc_c[i](c) + net = self.blocks[i](net) + out = self.fc_out(F.relu(net)) + + return xyz.jagged_like(out) diff --git a/surface_reconstruction/nksr/nksr_fvdb/nn/network.py b/surface_reconstruction/nksr/nksr_fvdb/nn/network.py new file mode 100644 index 0000000..7e51d09 --- /dev/null +++ b/surface_reconstruction/nksr/nksr_fvdb/nn/network.py @@ -0,0 +1,42 @@ +# Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property +# and proprietary rights in and to this software, related documentation +# and any modifications thereto. Any use, reproduction, disclosure or +# distribution of this software and related documentation without an express +# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited. + +import torch +import torch.nn as nn + +from nksr_fvdb.nn.encdec import MultiscalePointDecoder, PointEncoder +from nksr_fvdb.nn.unet import SparseStructureNet + + +class NKSRNetwork(nn.Module): + def __init__(self, hparams): + super().__init__() + self.hparams = hparams + self.encoder = PointEncoder(dim=6) + + self.sdf_decoder = MultiscalePointDecoder( + c_each_dim=self.hparams.kernel_dim, multiscale_depths=self.hparams.tree_depth, coords_depths=[2, 3] + ) + normal_channels = 3 + + # self.unet = SparseStructureNet( + # in_channels=32, + # num_blocks=self.hparams.tree_depth, + # basis_channels=self.hparams.kernel_dim, + # normal_channels=normal_channels, + # f_maps=self.hparams.unet.f_maps, + # udf_branch_dim=16 if self.hparams.udf.enabled else 0, + # ) + + if self.hparams.udf.enabled: + self.udf_decoder = MultiscalePointDecoder( + c_each_dim=16, + multiscale_depths=self.hparams.tree_depth, + out_init=5 * self.hparams.voxel_size, + coords_depths=[2, 3], + ) diff --git a/surface_reconstruction/nksr/nksr_fvdb/nn/unet.py b/surface_reconstruction/nksr/nksr_fvdb/nn/unet.py new file mode 100644 index 0000000..123117b --- /dev/null +++ b/surface_reconstruction/nksr/nksr_fvdb/nn/unet.py @@ -0,0 +1,280 @@ +# Copyright (c) 2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property +# and proprietary rights in and to this software, related documentation +# and any modifications thereto. Any use, reproduction, disclosure or +# distribution of this software and related documentation without an express +# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited. + + +import fvdb +import fvdb.nn as fvnn +import torch +import torch.nn as nn +from fvdb import GridBatch, JaggedTensor + +from nksr_fvdb.svh import SparseFeatureHierarchy, VoxelStatus + + +class SparseConvBlock(nn.Sequential): + def __init__(self, in_channels: int, out_channels: int, order: str, num_groups: int, kernel_size: int = 3): + super().__init__() + for i, char in enumerate(order): + if char == "r": + self.add_module("ReLU", fvnn.ReLU(inplace=True)) + elif char == "l": + self.add_module("LeakyReLU", fvnn.LeakyReLU(negative_slope=0.1, inplace=True)) + elif char == "c": + # add learnable bias only in the absence of batchnorm/groupnorm + self.add_module( + "Conv", + fvnn.SparseConv3d(in_channels, out_channels, kernel_size, 1, bias="g" not in order), + ) + elif char == "g": + if i < order.index("c"): + num_channels = in_channels + else: + num_channels = out_channels + + # use only one group if the given number of groups is greater than the number of channels + if num_channels < num_groups: + num_groups = 1 + + assert num_channels % num_groups == 0, ( + f"Expected number of channels in input to be divisible by num_groups. " + f"num_channels={num_channels}, num_groups={num_groups}" + ) + + self.add_module("GroupNorm", fvnn.GroupNorm(num_groups=num_groups, num_channels=num_channels)) + + else: + raise NotImplementedError + + +class SparseDoubleConv(nn.Sequential): + def __init__(self, in_channels: int, out_channels: int, order: str, num_groups: int, encoder: bool): + super().__init__() + if encoder: + conv1_in_channels = in_channels + conv1_out_channels = out_channels // 2 + if conv1_out_channels < in_channels: + conv1_out_channels = in_channels + conv2_in_channels, conv2_out_channels = conv1_out_channels, out_channels + else: + # we're in the decoder path, decrease the number of channels in the 1st convolution + conv1_in_channels, conv1_out_channels = in_channels, out_channels + conv2_in_channels, conv2_out_channels = out_channels, out_channels + + self.add_module("SingleConv1", SparseConvBlock(conv1_in_channels, conv1_out_channels, order, num_groups)) + self.add_module("SingleConv2", SparseConvBlock(conv2_in_channels, conv2_out_channels, order, num_groups)) + + +class SparseHead(nn.Sequential): + def __init__(self, in_channels: int, out_channels: int, order: str, num_groups: int, enhanced: bool = False): + super().__init__() + self.add_module("SingleConv", SparseConvBlock(in_channels, in_channels, order, num_groups)) + if enhanced: + mid_channels = min(64, in_channels) + self.add_module("SingleConv2", SparseConvBlock(in_channels, mid_channels, order, num_groups)) + self.add_module("OneConv0", SparseConvBlock(mid_channels, mid_channels, order, num_groups, kernel_size=1)) + self.add_module("OutConv", fvnn.SparseConv3d(mid_channels, out_channels, 1, bias=True)) + else: + self.add_module("OutConv", fvnn.SparseConv3d(in_channels, out_channels, 1, bias=True)) + + +class FeaturesSet: + def __init__(self): + self.structure_features = {} + self.normal_features = {} + self.basis_features = {} + self.udf_features = {} + + +class SparseStructureNet(nn.Module): + def __init__( + self, + in_channels: int, + num_blocks: int, + basis_channels: int, + normal_channels: int = 3, + f_maps: int = 64, + order: str = "gcr", + num_groups: int = 8, + neck_type: str = "dense", + neck_expand: int = 1, + udf_branch_dim: int = 16, + ): + super().__init__() + + n_features = [in_channels] + [f_maps * 2**k for k in range(num_blocks)] + self.num_blocks = num_blocks + self.neck_type = neck_type + self.neck_expand = neck_expand + self.basis_channels = basis_channels + self.normal_channels = normal_channels + self.udf_branch_dim = udf_branch_dim + + self.downsamplers = nn.ModuleList() + self.encoders = nn.ModuleList() + self.upsamplers = nn.ModuleList() + self.decoders = nn.ModuleList() + self.struct_heads = nn.ModuleList() + self.normal_heads = nn.ModuleList() + self.basis_heads = nn.ModuleList() + + if self.udf_branch_dim == 0: + self.udf_heads = [None for _ in range(num_blocks)] + else: + self.udf_heads = nn.ModuleList() + + for layer_idx in range(num_blocks): + self.encoders.add_module( + f"Enc{layer_idx}", + SparseDoubleConv(n_features[layer_idx], n_features[layer_idx + 1], order, num_groups, True), + ) + + for layer_idx in range(1, num_blocks): + self.downsamplers.add_module(f"Down{layer_idx}", fvnn.MaxPool(kernel_size=2)) + + for layer_idx in range(-1, -num_blocks - 1, -1): + self.struct_heads.add_module(f"Struct{layer_idx}", SparseHead(n_features[layer_idx], 3, order, num_groups)) + + if self.udf_branch_dim > 0: + self.udf_heads.add_module( + f"UDF{layer_idx}", SparseHead(n_features[layer_idx], self.udf_branch_dim, order, num_groups) + ) + if self.normal_channels == 0: + self.normal_heads.add_module(f"Normal{layer_idx}", None) + else: + self.normal_heads.add_module( + f"Normal{layer_idx}", SparseHead(n_features[layer_idx], self.normal_channels, order, num_groups) + ) + self.basis_heads.add_module( + f"Basis{layer_idx}", + SparseHead(n_features[layer_idx], self.basis_channels, order, num_groups, enhanced=True), + ) + + if layer_idx < -1: + self.decoders.add_module( + f"Dec{layer_idx}", + SparseDoubleConv( + n_features[layer_idx + 1] + n_features[layer_idx], + n_features[layer_idx], + order, + num_groups, + False, + ), + ) + up_module = fvnn.UpsamplingNearest(2) + self.upsamplers.add_module(f"Up{layer_idx}", up_module) + + self.padding = fvnn.FillToGrid() + + def build_neck_grid(self, sparse_grid: GridBatch, dense_svh: SparseFeatureHierarchy): + sparse_coords = sparse_grid.ijk + n_padding = (self.neck_expand - 1) // 2 + + if self.neck_type == "dense": + all_coords = [] + for b in range(sparse_grid.grid_count): + min_bound = torch.min(sparse_coords[b].jdata, dim=0).values.cpu().numpy() - n_padding + max_bound = torch.max(sparse_coords[b].jdata, dim=0).values.cpu().numpy() + 1 + n_padding + cx = torch.arange(min_bound[0], max_bound[0], dtype=torch.int32, device=sparse_coords.device) + cy = torch.arange(min_bound[1], max_bound[1], dtype=torch.int32, device=sparse_coords.device) + cz = torch.arange(min_bound[2], max_bound[2], dtype=torch.int32, device=sparse_coords.device) + coords = torch.stack(torch.meshgrid(cx, cy, cz, indexing="ij"), dim=3).view(-1, 3) + all_coords.append(coords) + all_coords = JaggedTensor(all_coords) + dense_svh.build_from_grid_coords(dense_svh.depth - 1, all_coords) + + else: + dense_svh.build_from_grid_coords( + dense_svh.depth - 1, + sparse_coords, + [-n_padding, -n_padding, -n_padding], + [n_padding, n_padding, n_padding], + ) + + def forward( + self, + feat: JaggedTensor, + encoder_svh: SparseFeatureHierarchy, + adaptive_depth: int, + ): + res = FeaturesSet() + feat_depth = 0 + vdb_tensor = VDBTensor(encoder_svh.grids[feat_depth], feat) + + # Down-sample + encoder_features = {} + for module, downsampler in zip(self.encoders, [None] + list(self.downsamplers)): + if downsampler is not None: + vdb_tensor = downsampler(vdb_tensor, ref_coarse_data=encoder_svh.grids[feat_depth + 1]) + feat_depth += 1 + vdb_tensor = module(vdb_tensor) + encoder_features[feat_depth] = vdb_tensor + + # Bottleneck processing + decoder_svh = SparseFeatureHierarchy(encoder_svh.voxel_size, encoder_svh.depth, encoder_svh.device) + decoder_tmp_svh = SparseFeatureHierarchy(encoder_svh.voxel_size, encoder_svh.depth, encoder_svh.device) + + self.build_neck_grid(encoder_svh.grids[feat_depth], decoder_tmp_svh) + vdb_tensor = self.padding(vdb_tensor, decoder_tmp_svh.grids[feat_depth]) + + # Up-sample + upsample_mask = None + for module, upsampler, struct_conv, normal_conv, basis_conv, udf_conv in zip( + [None] + list(self.decoders), + [None] + list(self.upsamplers), + self.struct_heads, + self.normal_heads, + self.basis_heads, + self.udf_heads, + ): + if module is not None: + vdb_tensor = upsampler(vdb_tensor, mask=upsample_mask) + feat_depth -= 1 + decoder_tmp_svh.build_from_grid(feat_depth, vdb_tensor.grid) + + enc_feat = self.padding(encoder_features[feat_depth], vdb_tensor) + vdb_tensor = VDBTensor.cat([enc_feat, vdb_tensor], dim=1) + + vdb_tensor = module(vdb_tensor) + + # Do structure inference. + res.structure_features[feat_depth] = struct_conv(vdb_tensor).feature + + if udf_conv is not None: + res.udf_features[feat_depth] = udf_conv(vdb_tensor).feature + + struct_decision = torch.argmax(res.structure_features[feat_depth].jdata, dim=1).byte() + + exist_mask = struct_decision != VoxelStatus.VS_NON_EXIST.value + + # If the predicted structure is empty, then stop early. + # (Related branch will not have gradient) + if not torch.any(exist_mask): + break + + dec_ijk = decoder_tmp_svh.grids[feat_depth].ijk.r_masked_select(exist_mask) + decoder_svh.build_from_grid_coords(feat_depth, dec_ijk) + + vdb_tensor = self.padding(vdb_tensor, decoder_svh.grids[feat_depth]) + upsample_mask = ( + decoder_svh.grids[feat_depth] + .fill_to_grid( + (struct_decision == VoxelStatus.VS_EXIST_CONTINUE.value).float(), decoder_tmp_svh.grids[feat_depth] + ) + .type(torch.bool) + ) + + # Do normal&basis prediction. + if feat_depth < adaptive_depth and normal_conv is not None: + res.normal_features[feat_depth] = normal_conv(vdb_tensor).feature + res.basis_features[feat_depth] = basis_conv(vdb_tensor).feature + + # If next level is for sure empty, then stop earlier + if not torch.any(upsample_mask.jdata): + break + + return res, decoder_svh, decoder_tmp_svh diff --git a/surface_reconstruction/nksr/nksr_fvdb/presets.py b/surface_reconstruction/nksr/nksr_fvdb/presets.py new file mode 100644 index 0000000..1d481bd --- /dev/null +++ b/surface_reconstruction/nksr/nksr_fvdb/presets.py @@ -0,0 +1,25 @@ +# Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property +# and proprietary rights in and to this software, related documentation +# and any modifications thereto. Any use, reproduction, disclosure or +# distribution of this software and related documentation without an express +# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited. + +from omegaconf import DictConfig + +all_presets = { + "ks-sdf": DictConfig( + { + "url": "https://huggingface.co/heiwang1997/nksr-checkpoints/resolve/main/checkpoints/ks.pth", + "voxel_size": 0.1, + "kernel_dim": 4, + "tree_depth": 4, + "adaptive_depth": 2, + "unet": {"f_maps": 32}, + "udf": {"enabled": True}, + "interpolator": {"n_hidden": 2, "hidden_dim": 16}, + "density_range": [1.0, 20.0], + } + ), +} diff --git a/surface_reconstruction/nksr/nksr_fvdb/reconstructor.py b/surface_reconstruction/nksr/nksr_fvdb/reconstructor.py new file mode 100644 index 0000000..464fe85 --- /dev/null +++ b/surface_reconstruction/nksr/nksr_fvdb/reconstructor.py @@ -0,0 +1,101 @@ +# Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property +# and proprietary rights in and to this software, related documentation +# and any modifications thereto. Any use, reproduction, disclosure or +# distribution of this software and related documentation without an express +# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited. + +import logging +import math + +import torch +import torch.nn as nn +from fvdb import JaggedTensor +from fvdb.nn.simple_unet import SimpleUNet + +from nksr_fvdb.fields import BaseField, LayerField, NeuralField +from nksr_fvdb.nn.network import NKSRNetwork +from nksr_fvdb.presets import all_presets +from nksr_fvdb.svh import SparseFeatureHierarchy + +logger = logging.getLogger(__name__) + + +class Reconstructor: + """ + Main Reconstructor class that reconstructs an implicit field from a point cloud. + """ + + def __init__(self, device: torch.device | str, preset: str = "ks"): + """ + Args: + device (torch.device): device to run the reconstructor on. + preset (str): name of the reconstructor preset. + """ + self.device = device + self.chunk_tmp_device = self.device + self.hparams = all_presets[preset] + self.network = NKSRNetwork(self.hparams).to(self.device).eval().requires_grad_(False) + # ckpt_data = torch.hub.load_state_dict_from_url(self.hparams.url) + # self.network.load_state_dict(ckpt_data["state_dict"]) + + self.simple_unet = ( + SimpleUNet( + in_channels=32, + base_channels=32, + out_channels=32, + channel_growth_rate=2, + ) + .to(self.device) + .eval() + .requires_grad_(False) + ) + + def reconstruct( + self, + xyz: torch.Tensor | JaggedTensor, + normal: torch.Tensor | JaggedTensor, + voxel_size: float, + ) -> BaseField: + """ + + Args: + xyz (torch.Tensor): (N, 3) input point positions + normal (torch.Tensor): (N, 3) input point normals + voxel_size (float): the voxel size of the input point cloud to use + + Returns: + field (Field): the implicit field to extract mesh from. + """ + if isinstance(xyz, torch.Tensor): + xyz = JaggedTensor.from_tensor(xyz) + + if isinstance(normal, torch.Tensor): + normal = JaggedTensor.from_tensor(normal) + + if (global_scale := voxel_size / self.hparams.voxel_size) != 1.0: + xyz = xyz / global_scale + + svh = SparseFeatureHierarchy( + voxel_size=self.hparams.voxel_size, depth=self.hparams.tree_depth, device=self.device + ) + svh.build_point_splatting(xyz) + feat = self.network.encoder(xyz, normal, svh, 0) + + # Choose between our customized implementation and simple unet: + # feat, svh, udf_svh = self.network.unet(feat, svh, adaptive_depth=self.hparams.adaptive_depth) + feat = self.simple_unet(feat, grid=svh.grids[0]) + + output_field = NeuralField(svh=svh, decoder=self.network.sdf_decoder, features=feat.basis_features) + + if self.hparams.udf.enabled: + mask_field = NeuralField(svh=udf_svh, decoder=self.network.udf_decoder, features=feat.udf_features) + mask_field.set_level_set(2 * self.hparams.voxel_size) + else: + mask_field = LayerField(svh, self.hparams.adaptive_depth) + + output_field.mask_field = mask_field + output_field.set_scale(global_scale) + + return output_field diff --git a/surface_reconstruction/nksr/nksr_fvdb/svh.py b/surface_reconstruction/nksr/nksr_fvdb/svh.py new file mode 100644 index 0000000..84bd3c4 --- /dev/null +++ b/surface_reconstruction/nksr/nksr_fvdb/svh.py @@ -0,0 +1,122 @@ +# Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property +# and proprietary rights in and to this software, related documentation +# and any modifications thereto. Any use, reproduction, disclosure or +# distribution of this software and related documentation without an express +# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited. + + +from enum import Enum + +import torch +from fvdb import GridBatch, JaggedTensor + + +class VoxelStatus(Enum): + """ + Status of a voxel in the structure pruning stage. + """ + + # The voxel should not exist (pruned) + VS_NON_EXIST = 0 + + # The voxel exists, but would not have children + VS_EXIST_STOP = 1 + + # The voxel exists, and would have children + VS_EXIST_CONTINUE = 2 + + +class SparseFeatureHierarchy: + """A hierarchy of grid batch, where voxel corners align with the origin""" + + def __init__(self, voxel_size: float, depth: int, device): + self.device = device + self.voxel_size = voxel_size + self.depth = depth + self.grids: list[GridBatch] = [ + GridBatch.from_zero_voxels( + self.device, + self.get_voxel_size(d), + self.get_origin(d), + ) + for d in range(self.depth) + ] + + def __repr__(self): + text = f"SparseFeatureHierarchy - {self.depth} layers, Voxel size = {self.voxel_size}" + text += "\n" + for d, d_grid in enumerate(self.grids): + text += f"\t[{d} {d_grid.num_voxels} voxels]" + return text + "\n" + + def get_voxel_size(self, depth: int) -> float: + return self.voxel_size * (2**depth) + + def get_origin(self, depth: int) -> float: + return 0.5 * self.voxel_size * (2**depth) + + def get_voxel_centers(self, depth: int) -> JaggedTensor: + grid = self.grids[depth] + return grid.voxel_to_world(grid.ijk.float()) + + def get_f_bound(self) -> tuple[JaggedTensor, JaggedTensor]: + grid = self.grids[self.depth - 1] + grid_coords = grid.ijk.float() + min_extent = grid.voxel_to_world(grid_coords.jmin(dim=0)[0] - 1.5) + max_extent = grid.voxel_to_world(grid_coords.jmax(dim=0)[0] + 1.5) + return min_extent, max_extent + + def evaluate_voxel_status(self, grid: GridBatch, depth: int): + """ + Evaluate the voxel status of given coordinates + :param grid: Featuregrid Grid + :param depth: int + :return: (N, ) byte tensor, with value 0,1,2 + """ + status = torch.full((grid.num_voxels,), VoxelStatus.VS_NON_EXIST.value, dtype=torch.uint8, device=self.device) + + if self.grids[depth] is not None: + exist_idx = grid.ijk_to_index(self.grids[depth].ijk) + status[exist_idx[exist_idx != -1]] = VoxelStatus.VS_EXIST_STOP.value + + if depth > 0 and self.grids[depth - 1] is not None: + child_coords = torch.div(self.grids[depth - 1].ijk, 2, rounding_mode="floor") + child_idx = grid.ijk_to_index(child_coords) + status[child_idx[child_idx != -1]] = VoxelStatus.VS_EXIST_CONTINUE.value + + return status + + def get_test_grid(self, depth: int = 0, resolution: int = 2): + grid = self.grids[depth] + assert grid is not None + primal_coords = grid.ijk.float() + box_coords = torch.linspace(-0.5, 0.5, resolution, device=self.device) + box_coords = torch.stack(torch.meshgrid(box_coords, box_coords, box_coords, indexing="ij"), dim=3) + box_coords = box_coords.view(-1, 3) + query_pos = primal_coords.unsqueeze(1) + box_coords.unsqueeze(0) + query_pos = grid.voxel_to_world(query_pos.view(-1, 3)) + return query_pos, primal_coords + + def to_(self, device: torch.device | str): + device = torch.device(device) + if device == self.device: + return + self.device = device + self.grids = [v.to(device) if v is not None else None for v in self.grids] + + def build_iterative_coarsening(self, pts: JaggedTensor): + assert pts.device == self.device, f"Device not match {pts.device} vs {self.device}." + self.grids[0] = GridBatch.from_points( + pts, voxel_sizes=self.get_voxel_size(0), origins=self.get_origin(0), device=self.device + ) + for d in range(1, self.depth): + self.grids[d] = self.grids[d - 1].coarsened_grid(2) + + def build_point_splatting(self, pts: JaggedTensor): + assert pts.device == self.device, f"Device not match {pts.device} vs {self.device}." + for d in range(self.depth): + self.grids[d] = GridBatch.from_nearest_voxels_to_points( + pts, voxel_sizes=self.get_voxel_size(d), origins=self.get_origin(d), device=self.device + ) diff --git a/surface_reconstruction/nksr/pyproject.toml b/surface_reconstruction/nksr/pyproject.toml new file mode 100644 index 0000000..261f8dc --- /dev/null +++ b/surface_reconstruction/nksr/pyproject.toml @@ -0,0 +1,30 @@ +[build-system] +requires = ["setuptools>=40.8.0", "wheel", "torch"] +build-backend = "setuptools.build_meta" + +[project] +name = "nksr-fvdb" +description = "Neural Kernel Surface Reconstruction (fVDB version)" +authors = [ + {name = "The NKSR Authors", email = "jiahuih@nvidia.com"}, +] +readme = "README.md" +requires-python = ">=3.10" +dynamic = ["version"] +dependencies = [ + "torch", + "omegaconf", +] + +[tool.setuptools.dynamic] +version = {attr = "nksr_fvdb.__version__"} + +[tool.ruff] +target-version = "py310" +line-length = 120 + +[tool.ruff.lint.isort] +known-first-party = ["nksr_fvdb"] +section-order = ["future", "standard-library", "third-party", "first-party", "local-folder"] +lines-between-types = 1 +lines-after-imports = 2