fps on cpu

Author: nicolas-chaulet

cpu/include/fps.h

@@ -0,0 +1,3 @@
+#pragma once
+#include <torch/extension.h>
+at::Tensor fps(at::Tensor points, const int nsamples, bool random = true);

cpu/src/bindings.cpp

@@ -1,5 +1,5 @@
 #include "ball_query.h"
-// #include "fps.h"
+#include "fps.h"
 #include "interpolate.h"
 #include "knn.h"
 
@@ -11,6 +11,7 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m)
     m.def("knn_interpolate", &knn_interpolate, "", "features"_a, "idx"_a, "weights"_a);
     m.def("knn_interpolate_grad", &knn_interpolate_grad, "", "grad_out"_a, "idx"_a, "weights"_a,
           "m"_a);
+    m.def("fps", &fps, "", "points"_a, "num_samples"_a, "random"_a);
 
     m.def("ball_query", &ball_query,
           "compute the radius search of a point cloud using nanoflann"

cpu/src/fps.cpp

@@ -0,0 +1,37 @@
+#include <torch/extension.h>
+
+#include "compat.h"
+#include "utils.h"
+
+at::Tensor get_dist(at::Tensor x, ptrdiff_t index)
+{
+    return (x - x[index]).norm(2, 1);
+}
+
+at::Tensor fps(at::Tensor points, const int nsamples, bool random)
+{
+    CHECK_CONTIGUOUS(points);
+
+    auto out_options = torch::TensorOptions().dtype(torch::kLong).device(torch::kCPU);
+    auto batch_size = points.size(0);
+    auto out = torch::empty({batch_size, nsamples}, out_options);
+    auto out_a = out.accessor<long, 2>();
+
+    for (ptrdiff_t b = 0; b < batch_size; b++)
+    {
+        auto y = points[b];
+        ptrdiff_t start = 0;
+        if (random)
+            start = at::randperm(y.size(0), out_options).DATA_PTR<int64_t>()[0];
+
+        out_a[b][0] = start;
+        auto dist = get_dist(y, start);
+        for (ptrdiff_t i = 1; i < nsamples; i++)
+        {
+            ptrdiff_t argmax = dist.argmax().DATA_PTR<int64_t>()[0];
+            out_a[b][i] = argmax;
+            dist = at::min(dist, get_dist(y, argmax));
+        }
+    }
+    return out;
+}

test/test_fps.py

@@ -0,0 +1,25 @@
+import unittest
+import torch
+import os
+import sys
+
+ROOT = os.path.join(os.path.dirname(os.path.realpath(__file__)), "..")
+sys.path.insert(0, ROOT)
+
+from torch_points.points_cpu import fps
+
+
+class TestFps(unittest.TestCase):
+    def test_simplecpu(self):
+        points = torch.tensor([[[0, 0, 0], [1, 0, 0], [2, 0, 0]], [[-1, 1, 0], [0, 0, 10], [0, 0, 2]]]).float()
+        idx = fps(points, 2, False)
+        torch.testing.assert_allclose(idx, torch.tensor([[0, 2], [0, 1]]))
+
+    def test_random(self):
+        points = torch.randn(10, 100, 3)
+        idx = fps(points, 2, True)
+        self.assertNotEqual(idx[0][0], 0)
+
+
+if __name__ == "__main__":
+    unittest.main()

torch_points/torchpoints.py

@@ -17,7 +17,7 @@ def forward(ctx, xyz, npoint):
         if xyz.is_cuda:
             return tpcuda.furthest_point_sampling(xyz, npoint)
         else:
-            raise NotImplementedError
+            return tpcpu.fps(xyz, npoint, True)
 
     @staticmethod
     def backward(xyz, a=None):