52 lines
1.4 KiB
Python
52 lines
1.4 KiB
Python
from abc import ABC
|
|
|
|
import numpy as np
|
|
|
|
|
|
class _Sampler(ABC):
|
|
|
|
def __init__(self, K, **kwargs):
|
|
self.k = K
|
|
self.kwargs = kwargs
|
|
|
|
def __call__(self, *args, **kwargs):
|
|
raise NotImplementedError
|
|
|
|
|
|
class RandomSampling(_Sampler):
|
|
|
|
def __init__(self, *args, **kwargs):
|
|
super(RandomSampling, self).__init__(*args, **kwargs)
|
|
|
|
def __call__(self, pts, *args, **kwargs):
|
|
rnd_indexs = np.random.choice(np.arange(pts.shape[0]), min(self.k, pts.shape[0]), replace=False)
|
|
return rnd_indexs
|
|
|
|
|
|
class FarthestpointSampling(_Sampler):
|
|
|
|
def __init__(self, *args, **kwargs):
|
|
super(FarthestpointSampling, self).__init__(*args, **kwargs)
|
|
|
|
@staticmethod
|
|
def calc_distances(p0, points):
|
|
return ((p0[:3] - points[:, :3]) ** 2).sum(axis=1)
|
|
|
|
def __call__(self, pts, *args, **kwargs):
|
|
|
|
if pts.shape[0] < self.k:
|
|
return pts
|
|
|
|
else:
|
|
farthest_pts = np.zeros((self.k, pts.shape[1]))
|
|
farthest_pts_idx = np.zeros(self.k, dtype=np.int)
|
|
farthest_pts[0] = pts[np.random.randint(len(pts))]
|
|
distances = self.calc_distances(farthest_pts[0], pts)
|
|
for i in range(1, self.k):
|
|
farthest_pts_idx[i] = np.argmax(distances)
|
|
farthest_pts[i] = pts[farthest_pts_idx[i]]
|
|
|
|
distances = np.minimum(distances, self.calc_distances(farthest_pts[i], pts))
|
|
|
|
return farthest_pts_idx
|