Home

Awesome

Neural Splines: Fitting 3D Surfaces with Inifinitely-Wide Neural Networks

Neural Splines Teaser This repository contains the official implementation of the CVPR 2021 (Oral) paper Neural Splines: Fitting 3D Surfaces with Infinitely-Wide Neural Networks.

Setup Instructions

System Requirements

Neural Splines uses FALKON, a state-of-the-art kernel ridge regression solver to fit surfaces on one or more GPUs. We thus require at least one GPU to run Neural Splines. We additionally require a working version of the CUDA compiler nvcc. We recommend running this code on a machine with a lot of memory if you want to reconstruct large point clouds since Neural Splines stores an MxM preconditioner matrix in CPU memory (where M is the number of Nystrom samples).

Installing Dependencies

Neural splines has several dependencies which must be installed before it can be used. Some of these dependencies must be built and take time to install. There are three ways to install dependencies:

Installing Dependencies with conda

Simply run

conda env create -f environment.yml

and then go grab a coffee ☕. When you get back, you will have a conda environment called neural-splines with the right dependencies installed.

Installing Dependencies with pip

We include several requirement-*.txt files in the requirements directory depending on your version of cuda. Choose the right file for your installation then run

pip install -r requirements/requirements-cuda<VERSION>.txt

and then go grab a coffee ☕.

Installing Dependencies Manually (Not Recommended)

You will need to install the following dependencies manually to use this repository:

You will also need to build the following dependencies from source. The easiest way to do this is with pip (see commands below), but you can also clone the linked repositories and run setup.py install:

Testing Your Installation

⚠️ WARNING ⚠️ Due to a bug in KeOps, the first time you use any code in this repository will throw a ModuleNotFoundError. All subsequent invocations of Neural Splines should work.

  1. Download and unzip the example point clouds here
  2. Unzip the file, in the directory of this repository, which should produce a directory named demo_data
  3. Run python fit.py demo_data/bunny.ply 10_000 128 On the first run this will fail (see above, just rerun it). On the second run it will compile some kernels and then produce a file called recon.ply which should be a reconstructed Stanford Bunny. The image below shows the input points and reconstruction for the bunny,
  4. Run python fit-grid.py demo_data/living_room_33_500_per_m2.ply 10_000 512 8 which will produce another recon.ply mesh, this time of a full room as shown below.
<p align="center"> <span> <img src="imgs/bunny.png" alt="A reconstructed Stanford Bunny" width="33%"> <img src="imgs/room.png" alt="A reconstruced living room" width="33%"> </span> </p>

Using Neural Splines from the Command Line

There are two scripts in this repository to fit surfaces from the command line:

Reconstructing a point cloud with fit.py

fit.py fits an input point cloud using a single Neural Spline. This approach works best for relatively small inputs which don't have too much geometric complexity. fit.py takes least the following arguments

fit.py <INPUT_POINT_CLOUD> <NUM_NYSTROM_SAMPLES> <GRID_SIZE>

where

Reconstructing very large point clouds with fit-grid.py

fit-grid.py fits an input point cloud in chunks using a different Neural Spline per chunk. This approach works well when the input point cloud is large or has a lot of geometric complexity. fit-grid.py takes the following required arguments

fit-grid.py <INPUT_POINT_CLOUD> <NUM_NYSTROM_SAMPLES> <GRID_SIZE> <CELLS_PER_AXIS>

where

Furthermore, fit-grid.py accepts the following optional arguments:

Additional arguments to fit.py and fit-grid.py

Additionally, both fit.py and fit-grid.py accept the following optional arguments which can alter the behavior and performance of the fitting process:

Trimming Reconstructed Meshes

Neural Splines can sometimes add surface sheets far away from input points, to remove these, we include a surface trimming script (similar to Poisson Surface Reconstruction), which trims mesh faces away from the input points. To trim a surface, simply run:

python trim-surface.py <INPUT_POINT_CLOUD> <RECONSTRUCTED_MESH> <GRID_SIZE> <DISTANCE_THRESHOLD> --out <OUT_FILE>

where:

Using Neural Splines in Python

Neural Splines can be used directly from within python by importing the neural_splines module in this repository.

To reconstruct a surface using Neural Splines, use the function neural_splines.fit_model_to_pointcloud. It returns a model object with the same API as Skikit-Learn. NOTE: neural_splines.fit_model_to_pointcloud can additionally accept other optional arguments. Run help(neural_splines.fit_model_to_pointcloud) for details.

from neural_splines import fit_model_to_point_cloud

# x is a point cloud stored in a torch tensor of shape [N, 3]
# n is a tensor of unit normals (one per point) of shape [N, 3]
# num_ny is the number of Nystrom samples to use 
# eps is the finite differencing coefficient (see documentation above)
model = fit_model_to_pointcloud(x, n, num_ny, eps)

# Evaluate the neural spline at a point p
p = torch.tensor([[0.5, 0.5, 0.5]]).to(x)
f_p = model.predict(p)

To evaluate a fitted Neural Spline on a grid of points, you can use the function neural_splines.eval_model_on_grid. NOTE: neural_splines.eval_model_on_grid can also accept other optional arguments, run help(neural_splines.eval_model_on_grid) for details.

from neural_splines import eval_model_on_grid

# Assume model is a Neural Spline fitted with fit_model_to_point_cloud

# Bounding box of the point cloud x represented as a tuple (origin, size)
bbox = x.min(0)[0], x.max(0)[0] - x.min(0)[0]
grid_res = torch.tensor([128, 128, 128]).to(torch.int32)

recon = eval_model_on_grid(model, bbox, voxel_grid_size)  # a [128, 128, 128] shaped tensor representing the neural spline evaluated on a grid.