PyTorch fpsample

PyTorch efficient farthest point sampling (FPS) implementation, adopted from fpsample.

Currently, this project is under heavy development and not ready for production use. If you want to make a contribution on implementing the GPU version, please feel free to contact me and make PRs.

[!NOTE] Since the PyTorch capsules the native multithread implementation, this project is expected to have a much better performance than the fpsample implementation.

Installation

# Install from github
pip install git+https://github.com/leonardodalinky/pytorch_fpsample

# Build locally
pip install .

Usage

import torch_fpsample

x = torch.rand(64, 2048, 3)
# random sample
sampled_points, indices = torch_fpsample.sample(x, 1024)
# random sample with specific tree height
sampled_points, indices = torch_fpsample.sample(x, 1024, h=5)
# random sample with start point index (int)
sampled_points, indices = torch_fpsample.sample(x, 1024, start_idx=0)

> sampled_points.size(), indices.size()
Size([64, 1024, 3]), Size([64, 1024])

[!WARNING] Note: The GPU version is not implemented yet. Only CPU mode is available.

Reference

Bucket-based farthest point sampling (QuickFPS) is proposed in the following paper. The implementation is based on the author's Repo (CPU & GPU).

@article{han2023quickfps,
  title={QuickFPS: Architecture and Algorithm Co-Design for Farthest Point Sampling in Large-Scale Point Clouds},
  author={Han, Meng and Wang, Liang and Xiao, Limin and Zhang, Hao and Zhang, Chenhao and Xu, Xiangrong and Zhu, Jianfeng},
  journal={IEEE Transactions on Computer-Aided Design of Integrated Circuits and Systems},
  year={2023},
  publisher={IEEE}
}

Thanks to the authors for their great works.