How to use in cpu

[leexusheng]

Hi, thank you for sharing the code. Run the entire pipeline on CPU failed Traceback (most recent call last): File "/workspace/NKSR/examples/recons_waymo.py", line 29, in field = reconstructor.reconstruct(
File "/root/miniconda3/envs/NKSR/lib/python3.10/site-packages/nksr/init.py", line 194, in reconstruct return self.reconstruct_by_chunk( File "/root/miniconda3/envs/NKSR/lib/python3.10/site-packages/nksr/init.py", line 342, in reconstruct_by_chunk chunk_field = self.reconstruct(xyz, normal=normal, sensor=sensor, **kwargs) File "/root/miniconda3/envs/NKSR/lib/python3.10/site-packages/nksr/init.py", line 208, in reconstruct res = preprocess_fn(xyz, normal, sensor) File "/root/miniconda3/envs/NKSR/lib/python3.10/site-packages/nksr/init.py", line 114, in func normal, sub_indices = utils.estimate_normals( File "/root/miniconda3/envs/NKSR/lib/python3.10/site-packages/nksr/utils.py", line 230, in estimate_normals knn_dist, knn_indices = ext.pcproc.nearest_neighbours(xyz, knn) RuntimeError: xyz must be a CUDA tensor

thanks!!!!

[heiwang1997]

Hi CPU supporting is currently under an experimental stage (it's currently slow as multithreading is not enabled). To perform normal estimation using CPU (and hence run the entire pipeline on CPU), please install point_cloud_utils via pip, and modify nksr/utils.py L246 as:

indices = torch.from_numpy(indices).to(xyz).long()

Then reconstruct with the following code:

import nksr
import torch

from pycg import vis, exp
from pathlib import Path
import numpy as np
from common import load_waymo_example, warning_on_low_memory

from nksr import utils

def normal_func(xyz: torch.Tensor, normal: torch.Tensor, sensor: torch.Tensor):
    assert normal is None, "normal already exists"
    assert sensor is not None, "please provide sensor positions for consistent orientations"

    normal, sub_indices = utils.estimate_normals(
        xyz, sensor,
        knn=64, drop_threshold_degrees=85.0, backend='pcu'
    )
    xyz = xyz[sub_indices]
    return xyz, normal, None

if __name__ == '__main__':
    warning_on_low_memory(20000.0)
    xyz_np, sensor_np = load_waymo_example()

    device = torch.device("cpu")
    reconstructor = nksr.Reconstructor(device)
    reconstructor.chunk_tmp_device = torch.device("cpu")

    input_xyz = torch.from_numpy(xyz_np).float().to(device)
    input_sensor = torch.from_numpy(sensor_np).float().to(device)

    field = reconstructor.reconstruct(
        input_xyz, sensor=input_sensor, detail_level=None,
        # Minor configs for better efficiency (not necessary)
        approx_kernel_grad=True, solver_tol=1e-4, fused_mode=True, 
        # Chunked reconstruction (if OOM)
        # chunk_size=51.2,
        preprocess_fn=normal_func
    )

    # (Optional) Convert to CPU for mesh extraction
    # field.to_("cpu")
    # reconstructor.network.to("cpu")

    mesh = field.extract_dual_mesh(mise_iter=1)
    mesh = vis.mesh(mesh.v, mesh.f)

    vis.show_3d([mesh], [vis.pointcloud(xyz_np)])

[heiwang1997]

Added a CPU version! Please refer to this example. Thanks