Closed chenzhaiyu closed 1 year ago
You're right, this can cause errors. I think this is mostly a theoretical problem since point clouds of <300 will most likely be useless anyway. I also experimented with ball queries (worse results than kNN) where this kind of problem is much more common. Feel free to use this code I wrote for follow-up work (call to get_patch_points_knn instead of get_patch_points):
def model_space_to_patch_space(
pts_to_convert_ms: np.array, pts_patch_center_ms: np.array, patch_radius_ms: typing.Union[float, np.ndarray]):
pts_patch_center_ms_repeated = \
np.repeat(np.expand_dims(pts_patch_center_ms, axis=0), pts_to_convert_ms.shape[-2], axis=-2)
pts_patch_space = pts_to_convert_ms - pts_patch_center_ms_repeated
pts_patch_space = pts_patch_space / patch_radius_ms
return pts_patch_space
def cartesian_dist(vec_x: np.array, vec_y: np.array, axis=1) -> np.ndarray:
"""
L2 distance
:param vec_x: array[n, d]
:param vec_y: array[n, d]
:param axis: int
:return: array[n]
"""
dist = np.linalg.norm(vec_x - vec_y, axis=axis)
return dist
def cartesian_dist_1_n(vec_x: np.array, vec_y: np.array, axis=1) -> np.ndarray:
"""
L2 distance
:param vec_x: array[d]
:param vec_y: array[n, d]
:param axis: int
:return: array[n]
"""
vec_x_bc = np.tile(np.expand_dims(vec_x, 0), (vec_y.shape[0], 1))
dist = np.linalg.norm(vec_x_bc - vec_y, axis=axis)
return dist
def get_patch_radii(pts_patch: np.array, query_pts: np.array):
if pts_patch.shape[0] == 0:
patch_radius = 0.0
elif pts_patch.shape == query_pts.shape:
patch_radius = np.linalg.norm(pts_patch - query_pts, axis=0)
else:
dist = cartesian_dist(np.repeat(np.expand_dims(query_pts, axis=0), pts_patch.shape[0], axis=0),
pts_patch, axis=1)
patch_radius = np.max(dist, axis=-1)
return patch_radius
def get_patch_kdtree(
kdtree: spatial.KDTree, rng: np.random.RandomState,
query_point, patch_radius, points_per_patch):
if patch_radius <= 0.0:
pts_dists_ms, patch_pts_ids = kdtree.query(x=query_point, k=points_per_patch)
else:
patch_pts_ids = kdtree.query_ball_point(x=query_point, r=patch_radius)
pts_dists_ms = cartesian_dist_1_n(query_point, kdtree.data[patch_pts_ids])
patch_pts_ids = np.array(patch_pts_ids, dtype=np.int32)
point_count = patch_pts_ids.shape[0]
# if there are too many neighbors, pick a random subset
if point_count > points_per_patch:
random_pts_idx = rng.choice(np.arange(point_count), points_per_patch, replace=False)
patch_pts_ids = patch_pts_ids[random_pts_idx]
pts_dists_ms = pts_dists_ms[random_pts_idx]
# pad with zeros (ball query)
if point_count < points_per_patch:
missing_points = points_per_patch - point_count
padding = np.full((missing_points), -1, dtype=np.int32)
dist_padding = np.full((missing_points), 0, dtype=np.float32)
if point_count == 0:
patch_pts_ids = padding
pts_dists_ms = dist_padding
else:
patch_pts_ids = np.concatenate((patch_pts_ids, padding), axis=0)
pts_dists_ms = np.concatenate((pts_dists_ms, dist_padding), axis=0)
# pad with zeros (knn)
patch_pts_ids[patch_pts_ids == kdtree.n] = -1
return patch_pts_ids, pts_dists_ms
def get_patch_points_knn(pts_ms: np.ndarray, kdtree: scipy.spatial.KDTree, normals_ms: np.ndarray,
query_point, patch_radius, points_per_patch, rng, every_nth=1):
patch_pts_ids, pts_dists_ms = get_patch_kdtree(
kdtree=kdtree, rng=rng, query_point=query_point,
patch_radius=patch_radius,
points_per_patch=points_per_patch * every_nth)
patch_pts_ids = patch_pts_ids[::every_nth]
# find -1 ids for padding
patch_pts_pad_ids = patch_pts_ids == -1
patch_pts_ids[patch_pts_pad_ids] = 0
pts_patch_ms = pts_ms[patch_pts_ids, :]
normals_patch_ms = normals_ms[patch_pts_ids, :]
# replace padding points with query point so that they appear in the patch origin
pts_patch_ms[patch_pts_pad_ids, :] = query_point
patch_radius_ms = get_patch_radii(pts_patch_ms, query_point) \
if patch_radius <= 0.0 else patch_radius
pts_patch_ps = model_space_to_patch_space(
pts_to_convert_ms=pts_patch_ms, pts_patch_center_ms=query_point,
patch_radius_ms=patch_radius_ms)
return patch_pts_ids, pts_patch_ps, pts_patch_ms, normals_patch_ms, patch_radius_ms, pts_dists_ms
Thanks :)
Hi, I bumped into an IndexError while training on my own point clouds:
Then I tried to identify a potential bug: When the point cloud is of less amount (e.g., 40) than the specified number of points per local patch (e.g., 300 by default). Then
KDTree.query
would returnself.n
(40 in this case) to indicate missing neighbors. These missing neighbors are not properly addressed later, regardless of padding -1 values, which results in the IndexError.A quick fix would be to change this line from
to
I didn't dig deeper on this but assume this would be a fix, hence this issue for your attention. What do you think?