def build_edges(
query, database, indices=None, r_max=1.0, k_max=10, return_indices=False, remove_self_loops=True
):
dists, idxs, nn, grid = frnn.frnn_grid_points(
points1=query.unsqueeze(0),
points2=database.unsqueeze(0),
lengths1=None,
lengths2=None,
K=k_max,
r=r_max,
grid=None,
return_nn=False,
return_sorted=True,
)
idxs = idxs.squeeze().int()
ind = torch.Tensor.repeat(
torch.arange(idxs.shape[0], device=device), (idxs.shape[1], 1), 1
).T.int()
positive_idxs = idxs >= 0
edge_list = torch.stack([ind[positive_idxs], idxs[positive_idxs]]).long()
# Reset indices subset to correct global index
if indices is not None:
edge_list[0] = indices[edge_list[0]]
# Remove self-loops
if remove_self_loops:
edge_list = edge_list[:, edge_list[0] != edge_list[1]]
if return_indices:
return edge_list, dists, idxs, ind
else:
return edge_list
def get_laplacian_weights(points,
normals,
iso_points,
iso_normals,
neighborhood_size=8):
"""
compute distance based on iso local neighborhood
"""
with autograd.no_grad():
P, _ = points.view(-1, 3).shape
search_radius = 0.15
dim = iso_points.view(-1, 3).norm(dim=-1).max() * 2
avg_spacing = iso_points.shape[1] / dim / 16
dists, idxs, nn, _ = frnn.frnn_grid_points(points,
iso_points,
K=1,
return_nn=True,
grid=None,
r=search_radius)
nn_normals = frnn.frnn_gather(iso_normals, idxs)
dists = torch.sum((points - nn.view_as(points)) *
(normals + nn_normals.view_as(normals)),
dim=-1)
dists = dists * dists
spatial_w = torch.exp(-dists * avg_spacing)
spatial_w[idxs.view_as(spatial_w) < 0] = 0
return spatial_w.view(points.shape[:-1])
def save_intermediate_results(fnames):
num_pcs = 1
k = 8
for fname in fnames:
# Problem here: our algorithm is problematic when
# there is only one grid in one dimension
# the drill is a long thin object which brings some trouble
if "drill_shaft" in fname:
continue
pc1 = torch.FloatTensor(
read_ply(fname)[None, :, :3]).cuda() # no need for normals
pc2 = torch.FloatTensor(
read_ply(fname)[None, :, :3]).cuda() # no need for normals
torch.save(pc1, "data/pc/" + fname.split('/')[-1][:-4] + '.pt')
print(fname)
print(pc1.shape)
num_points = pc1.shape[1]
pc1 = normalize_pc(pc1)
print(pc1.min(dim=1)[0], pc1.max(dim=1)[0])
pc2 = normalize_pc(pc2)
lengths1 = torch.ones(
(num_pcs, ), dtype=torch.long).cuda() * num_points
lengths2 = torch.ones(
(num_pcs, ), dtype=torch.long).cuda() * num_points
dists, idxs, nn, grid = frnn.frnn_grid_points(
pc1,
pc2,
lengths1,
lengths2,
K=k,
r=0.1,
filename=fname.split('/')[-1])
print(fname + " done!")
def get_iso_bilateral_weights(points, normals, iso_points, iso_normals):
""" find closest iso point, compute bilateral weight """
search_radius = 0.1
dim = iso_points.view(-1, 3).norm(dim=-1).max() * 2
avg_spacing = iso_points.shape[1] / dim / 16
dists, idxs, nn, _ = frnn.frnn_grid_points(points,
iso_points,
K=1,
return_nn=True,
grid=None,
r=search_radius)
iso_normals = F.normalize(iso_normals, dim=-1)
iso_normals = frnn.frnn_gather(iso_normals, idxs).view(1, -1, 3)
dists = torch.sum((nn.view_as(points) - points) * iso_normals, dim=-1)**2
# dists[idxs<0] = 10 * search_radius **2
# dists = dists.squeeze(-1)
spatial_w = torch.exp(-dists * avg_spacing)
normals = F.normalize(normals, dim=-1)
normal_w = torch.exp(-((1 - torch.sum(normals * iso_normals, dim=-1)) /
(1 - np.cos(np.deg2rad(60))))**2)
weight = spatial_w * normal_w
weight[idxs.view_as(weight) < 0] = 0
if not valid_value_mask(weight).all():
print("Illegal weights")
breakpoint()
return weight
def resample_uniformly(
pointclouds: Union[Pointclouds, torch.Tensor],
neighborhood_size: int = 8,
knn=None,
normals=None,
shrink_ratio: float = 0.5,
repulsion_mu: float = 1.0
) -> Union[Pointclouds, Tuple[torch.Tensor, torch.Tensor]]:
"""
resample first use wlop to consolidate point clouds to a smaller point clouds (halve the points)
then upsample with ear
Returns:
Pointclouds or padded points and number of points per batch
"""
import math
import frnn
points_init, num_points = convert_pointclouds_to_tensor(pointclouds)
batch_size = num_points.shape[0]
diag = (points_init.view(-1, 3).max(dim=0).values -
points_init.view(-1, 3).min(0).values).norm().item()
avg_spacing = math.sqrt(diag / points_init.shape[1])
search_radius = min(4 * avg_spacing * neighborhood_size, 0.2)
if knn is None:
dists, idxs, _, grid = frnn.frnn_grid_points(points_init,
points_init,
num_points,
num_points,
K=neighborhood_size + 1,
r=search_radius,
grid=None,
return_nn=False)
knn = _KNN(dists=dists[..., 1:], idx=idxs[..., 1:], knn=None)
# estimate normals
if isinstance(pointclouds, torch.Tensor):
normals = normals
else:
normals = pointclouds.normals_padded()
if normals is None:
normals = estimate_pointcloud_normals(
points_init,
neighborhood_size=neighborhood_size,
disambiguate_directions=False)
else:
normals = F.normalize(normals, dim=-1)
points = points_init
wlop_result = wlop(pointclouds,
ratio=shrink_ratio,
repulsion_mu=repulsion_mu)
up_result = upsample(wlop_result, num_points)
if is_pointclouds(pointclouds):
return up_result
return up_result.points_padded(), up_result.num_points_per_cloud()
def build_edges(query,
database,
indices=None,
r_max=1.0,
k_max=10,
return_indices=False):
"""
NOTE: These KNN/FRNN algorithms return the distances**2. Therefore we need to be careful when comparing them to the target distances (r_val, r_test), and to the margin parameter (which is L1 distance)
"""
if FRNN_AVAILABLE:
Dsq, I, nn, grid = frnn.frnn_grid_points(
points1=query.unsqueeze(0),
points2=database.unsqueeze(0),
lengths1=None,
lengths2=None,
K=k_max,
r=r_max,
grid=None,
return_nn=False,
return_sorted=True,
)
I = I.squeeze().int()
ind = torch.Tensor.repeat(torch.arange(I.shape[0], device=device),
(I.shape[1], 1), 1).T.int()
positive_idxs = I >= 0
edge_list = torch.stack([ind[positive_idxs], I[positive_idxs]]).long()
else:
if device == "cuda":
res = faiss.StandardGpuResources()
Dsq, I = faiss.knn_gpu(res=res, xq=query, xb=database, k=k_max)
elif device == "cpu":
index = faiss.IndexFlatL2(database.shape[1])
index.add(database)
Dsq, I = index.search(query, k_max)
ind = torch.Tensor.repeat(torch.arange(I.shape[0], device=device),
(I.shape[1], 1), 1).T.int()
edge_list = torch.stack([ind[Dsq <= r_max**2], I[Dsq <= r_max**2]])
# Reset indices subset to correct global index
if indices is not None:
edge_list[0] = indices[edge_list[0]]
# Remove self-loops
edge_list = edge_list[:, edge_list[0] != edge_list[1]]
if return_indices:
return edge_list, Dsq, I, ind
else:
return edge_list
def get_heat_kernel_weights(points,
normals,
iso_points,
iso_normals,
neighborhood_size=8,
sigma_p=0.4,
sigma_n=0.7):
"""
find closest k points, compute point2face distance, and normal distance
"""
P, _ = points.view(-1, 3).shape
search_radius = 0.15
dim = iso_points.view(-1, 3).norm(dim=-1).max()
avg_spacing = iso_points.shape[1] / (dim * 2**2) / 16
dists, idxs, nn, _ = frnn.frnn_grid_points(points,
iso_points,
K=neighborhood_size,
return_nn=True,
grid=None,
r=search_radius)
# features
with autograd.no_grad():
# normalize just to be sure
iso_normals = F.normalize(iso_normals, dim=-1, eps=1e-15)
normals = F.normalize(normals, dim=-1, eps=1e-15)
# features are composite of points and normals
features = torch.cat([points / sigma_p, normals / sigma_n], dim=-1)
features_iso = torch.cat([iso_points / sigma_p, iso_normals / sigma_n],
dim=-1)
# compute kernels (N,P,K) k(x,xi), xi \in Neighbor(x)
knn_idx = idxs
# features_nb = knn_gather(features_iso, knn_idx)
features_nb = frnn.frnn_gather(features_iso, knn_idx)
# (N,P,K,D)
features_diff = features.unsqueeze(2) - features_nb
features_dist = torch.sum(features_diff**2, dim=-1)
kernels = torch.exp(-features_dist)
kernels[knn_idx < 0] = 0
# N,P,K,K,D
features_diff_ij = features_nb[:, :, :,
None, :] - features_nb[:, :, None, :, :]
features_dist_ij = torch.sum(features_diff_ij**2, dim=-1)
kernel_matrices = torch.exp(-features_dist_ij)
kernel_matrices[knn_idx < 0] = 0
kernel_matrices[knn_idx.unsqueeze(-2).expand_as(kernel_matrices) < 0]
kernel_matrices_inv = pinverse(kernel_matrices)
weight = kernels.unsqueeze(
-2) @ kernel_matrices_inv @ kernels.unsqueeze(-1)
weight.clamp_max_(1.0)
return weight.view(points.shape[:-1])
def _compute_global_Vrk(self, pointclouds, refresh=True, **kwargs):
"""
determine variance scaler used in globally (see _compute_isotropic_Vrk)
Args:
pointclouds: pointclouds in object coorindates
Returns:
h_k: scaler
S_k: local frame
"""
if not refresh and self._Vrk_h is not None:
h_k = self._Vrk_h
else:
# compute average density
with torch.autograd.enable_grad():
pts_world = pointclouds.points_padded()
num_points_per_cloud = pointclouds.num_points_per_cloud()
if self.frnn_radius <= 0:
# use knn here
# logger_py.info("vrk knn points")
sq_dist, _, _ = ops3d.knn_points(pts_world,
pts_world,
num_points_per_cloud,
num_points_per_cloud,
K=7)
else:
sq_dist, _, _, _ = frnn.frnn_grid_points(pts_world,
pts_world,
num_points_per_cloud,
num_points_per_cloud,
K=7,
r=self.frnn_radius)
# logger_py.info("frnn and knn dist close: {}".format(torch.allclose(sq_dist, sq_dist2)))
sq_dist = sq_dist[:, :, 1:]
# knn search is unreliable, set sq_dist manually
sq_dist[num_points_per_cloud < 7] = 1e-3
h_k = 0.5 * sq_dist.max(dim=-1, keepdim=True)[0]
# prevent some outlier rendered be too large, or too small
h_k = h_k.mean(dim=1, keepdim=True).clamp(5e-5, 1e-3)
Vrk_h = gather_batch_to_packed(h_k,
pointclouds.packed_to_cloud_idx())
# Sk, a transformation from 2D local surface frame to 3D world frame
# Because isometry, two axis are equivalent, we can simply
# find two 3d vectors perpendicular to the point normals
# (totalP, 2, 3)
with torch.autograd.enable_grad():
normals = pointclouds.normals_packed()
u0 = F.normalize(torch.cross(normals,
normals + torch.rand_like(normals)),
dim=-1)
u1 = F.normalize(torch.cross(normals, u0), dim=-1)
Sk = torch.stack([u0, u1], dim=1)
Vrk = Vrk_h.view(-1, 1, 1) * Sk.transpose(1, 2) @ Sk
return Vrk, Sk
def frnn_grid_reuse(self):
dists, idxs, nn, _ = frnn.frnn_grid_points(self.pc1_frnn_reuse,
self.pc2_frnn_reuse,
self.lengths1_cuda,
self.lengths2_cuda,
K=self.K,
r=self.r,
grid=self.grid,
return_nn=True,
return_sorted=True)
return dists, idxs, nn
def output():
dists, idxs, nn, grid = frnn.frnn_grid_points(
points1,
points2,
lengths1,
lengths2,
K,
r,
grid=None,
return_nn=False,
return_sorted=True,
radius_cell_ratio=ratio)
torch.cuda.synchronize()
def frnn_grid(self):
dists, idxs, nn, grid = frnn.frnn_grid_points(self.pc1_frnn,
self.pc2_frnn,
self.lengths1_cuda,
self.lengths2_cuda,
K=self.K,
r=self.r,
grid=None,
return_nn=True,
return_sorted=True)
if self.grid is None:
self.grid = grid
return dists, idxs, nn
def frnn_2d(self):
dists, idxs, nn, grid = frnn.frnn_grid_points(
self.pc1,
self.pc2,
self.lengths1_cuda,
self.lengths2_cuda,
self.K,
self.r,
radius_cell_ratio=1.0
)
sorted_points2 = grid.sorted_points2
sorted_points2_idxs = grid.sorted_points2_idxs[:, :, None].long().expand(-1, -1, 2)
idxs_pc2 = torch.gather(self.pc2, 1, sorted_points2_idxs)
print(torch.allclose(sorted_points2, idxs_pc2))
def denoise_normals(points,
normals,
num_points,
sharpness_sigma=30,
knn_result=None,
neighborhood_size=16):
"""
Weights exp(-(1-<n, n_i>)/(1-cos(sharpness_sigma))), for i in a local neighborhood
"""
points, num_points = convert_pointclouds_to_tensor(points)
normals = F.normalize(normals, dim=-1)
if knn_result is None:
diag = (points.max(dim=-2)[0] - points.min(dim=-2)[0]).norm(dim=-1)
avg_spacing = math.sqrt(diag / points.shape[1])
search_radius = min(4 * avg_spacing * neighborhood_size, 0.2)
dists, idxs, _, grid = frnn.frnn_grid_points(points,
points,
num_points,
num_points,
K=neighborhood_size + 1,
r=search_radius,
grid=None,
return_nn=True)
knn_result = _KNN(dists=dists[..., 1:], idx=idxs[..., 1:], knn=None)
if knn_result.knn is None:
knn = frnn.frnn_gather(points, knn_result.idx, num_points)
knn_result = _KNN(idx=knn_result.idx, knn=knn, dists=knn_result.dists)
# filter out
knn_normals = frnn.frnn_gather(normals, knn_result.idx, num_points)
# knn_normals = frnn.frnn_gather(normals, self._knn_idx, num_points)
weights_n = (
(1 - torch.sum(knn_normals * normals[:, :, None, :], dim=-1)) /
sharpness_sigma)**2
weights_n = torch.exp(-weights_n)
inv_sigma_spatial = num_points / 2.0
spatial_dist = 16 / inv_sigma_spatial
deltap = knn - points[:, :, None, :]
deltap = torch.sum(deltap * deltap, dim=-1)
weights_p = torch.exp(-deltap * inv_sigma_spatial)
weights_p[deltap > spatial_dist] = 0
weights = weights_p * weights_n
# weights[self._knn_idx < 0] = 0
normals_denoised = torch.sum(knn_normals * weights[:, :, :, None], dim=-2) / \
eps_denom(torch.sum(weights, dim=-1, keepdim=True))
normals_denoised = F.normalize(normals_denoised, dim=-1)
return normals_denoised.view_as(normals)
def frnn_r(fname, N, K, r):
ragged = False
print(fname, N, K, r)
points1 = torch.load(fname)
points2 = torch.load(fname)
# print(points1.min(dim=1)[0], points1.max(dim=1)[0])
if N > 1:
points1 = points1.repeat(N, 1, 1)
points2 = points2.repeat(N, 1, 1)
if ragged:
lengths1 = torch.randint(low=1,
high=points1.shape[1],
size=(N, ),
dtype=torch.long,
device=points1.device)
lengths2 = torch.randint(low=1,
high=points2.shape[1],
size=(N, ),
dtype=torch.long,
device=points2.device)
else:
lengths1 = torch.ones(
(N, ), dtype=torch.long, device=points1.device) * points1.shape[1]
lengths2 = torch.ones(
(N, ), dtype=torch.long, device=points2.device) * points2.shape[1]
dists, idxs, nn, grid = frnn.frnn_grid_points(points1,
points2,
lengths1,
lengths2,
K,
r,
grid=None,
return_nn=False,
return_sorted=True)
return
import frnn
import torch
if __name__ == '__main__':
device = torch.device('cuda')
points = torch.load('pts.pth')[None, :, :].to(device).float()
n_points = torch.tensor([points.size(1)]).to(device).long()
K = 10
radius = 0.05
print(points.size(), n_points)
print(points.max(axis=1))
print(points.min(axis=1))
_, idxs, _, _ = frnn.frnn_grid_points(points,
points,
n_points,
n_points,
K,
radius,
grid=None,
return_nn=False,
return_sorted=False,
radius_cell_ratio=2)
def _compute_isotropic_Vrk(self, pointclouds, refresh=True, **kwargs):
"""
determine the variance in the local surface frame h * Sk.T @ Sk,
where Sk is 2x3 local surface coordinate to world coordinate.
determine the h_k in V_k^r = h_k*Id using nearest neighbor
heuristically h_k = mean(dist between points in a small neighbor)
The larger h_k is, the larger the splat is
NOTE: h_k in inverse to the definition in the paper, the larger h_k, the
larger the splats
Args:
pointclouds: pointcloud in object coordinate
Returns:
h_k: [N,3,3] tensor for each point
S_k: [N,2,3] local frame
"""
if not refresh and self._Vrk_h is not None and \
pointclouds.num_points_per_cloud().sum() == self._Vrk_h.shape[0]:
pass
else:
with torch.autograd.enable_grad():
pts_world = pointclouds.points_padded()
num_points_per_cloud = pointclouds.num_points_per_cloud()
if self.frnn_radius <= 0:
# logger_py.info("vrk knn points")
sq_dist, _, _ = ops3d.knn_points(pts_world,
pts_world,
num_points_per_cloud,
num_points_per_cloud,
K=7)
else:
sq_dist, _, _, _ = frnn.frnn_grid_points(pts_world,
pts_world,
num_points_per_cloud,
num_points_per_cloud,
K=7,
r=self.frnn_radius)
sq_dist = sq_dist[:, :, 1:]
# knn search is unreliable, set sq_dist manually
sq_dist[num_points_per_cloud < 7] = 1e-3
# (totalP, knnK)
sq_dist = ops3d.padded_to_packed(
sq_dist, pointclouds.cloud_to_packed_first_idx(),
num_points_per_cloud.sum().item())
# [totalP, ]
h_k = 0.5 * sq_dist.max(dim=-1, keepdim=True)[0]
# prevent some outlier rendered be too large, or too small
self._Vrk_h = h_k.clamp(5e-5, 0.01)
# Sk, a transformation from 2D local surface frame to 3D world frame
# Because isometry, two axis are equivalent, we can simply
# find two 3d vectors perpendicular to the point normals
# (totalP, 2, 3)
with torch.autograd.enable_grad():
normals = pointclouds.normals_packed()
u0 = F.normalize(torch.cross(normals,
normals + torch.rand_like(normals)),
dim=-1)
u1 = F.normalize(torch.cross(normals, u0), dim=-1)
Sk = torch.stack([u0, u1], dim=1)
Vrk = self._Vrk_h.view(-1, 1, 1) * Sk.transpose(1, 2) @ Sk
return Vrk, Sk
def wlop(pointclouds: PointClouds3D,
ratio: float = 0.5,
neighborhood_size=16,
iters=3,
repulsion_mu=0.5) -> PointClouds3D:
"""
Consolidation of Unorganized Point Clouds for Surface Reconstruction
Args:
pointclouds containing max J points per cloud
ratio: downsampling ratio (0, 1]
"""
P, num_points_P = convert_pointclouds_to_tensor(pointclouds)
# (N, 3, 2)
bbox = pointclouds.get_bounding_boxes()
# (N,)
diag = torch.norm(bbox[..., 0] - bbox[..., 1], dim=-1)
h = 4 * torch.sqrt(diag / num_points_P.float())
search_radius = min(h * neighborhood_size, 0.2)
theta_sigma_inv = 16 / h / h
if ratio < 1.0:
X0 = farthest_sampling(pointclouds, ratio=ratio)
elif ratio == 1.0:
X0 = pointclouds.clone()
else:
raise ValueError('ratio must be less or equal to 1.0')
# slightly perturb so that we don't find the same point when searching NN XtoP
offset = torch.randn_like(X0.points_packed()) * h * 0.1
X0.offset_(offset)
X, num_points_X = convert_pointclouds_to_tensor(X0)
def theta(r2):
return torch.exp(-r2 * theta_sigma_inv)
def eta(r):
return -r
def deta(r):
return torch.ones_like(r)
grid = None
dists, idxs, _, grid = frnn.frnn_grid_points(P,
P,
num_points_P,
num_points_P,
K=neighborhood_size + 1,
r=search_radius,
grid=grid,
return_nn=False)
knn_PtoP = _KNN(dists=dists[..., 1:], idx=idxs[..., 1:], knn=None)
deltapp = torch.norm(P.unsqueeze(-2) -
frnn.frnn_gather(P, knn_PtoP.idx, num_points_P),
dim=-1)
theta_pp_nn = theta(deltapp**2) # (B, P, K)
theta_pp_nn[knn_PtoP.idx < 0] = 0
density_P = torch.sum(theta_pp_nn, dim=-1) + 1
for it in range(iters):
# from each x find closest neighbors in pointclouds
dists, idxs, _, grid = frnn.frnn_grid_points(X,
P,
num_points_X,
num_points_P,
K=neighborhood_size,
r=search_radius,
grid=grid,
return_nn=False)
knn_XtoP = _KNN(dists=dists, idx=idxs, knn=None)
dists, idxs, _, _ = frnn.frnn_grid_points(X,
X,
num_points_X,
num_points_X,
K=neighborhood_size + 1,
r=search_radius,
grid=None,
return_nn=False)
knn_XtoX = _KNN(dists=dists[..., 1:], idx=idxs[..., 1:], knn=None)
# LOP local optimal projection
nn_XtoP = frnn.frnn_gather(P, knn_XtoP.idx, num_points_P)
epsilon = X.unsqueeze(-2) - frnn.frnn_gather(P, knn_XtoP.idx,
num_points_P)
delta = X.unsqueeze(-2) - frnn.frnn_gather(X, knn_XtoX.idx,
num_points_X)
# (B, I, I)
deltaxx2 = (delta**2).sum(dim=-1)
# (B, I, K)
deltaxp2 = (epsilon**2).sum(dim=-1)
# (B, I, K)
alpha = theta(deltaxp2) / eps_denom(epsilon.norm(dim=-1))
# (B, I, K)
beta = theta(deltaxx2) * deta(delta.norm(dim=-1)) / eps_denom(
delta.norm(dim=-1))
density_X = torch.sum(theta(deltaxx2), dim=-1) + 1
new_alpha = alpha / frnn.frnn_gather(
density_P.unsqueeze(-1), knn_XtoP.idx, num_points_P).squeeze(-1)
new_alpha[knn_XtoP.idx < 0] = 0
new_beta = density_X.unsqueeze(-1) * beta
new_beta[knn_XtoX.idx < 0] = 0
term_data = torch.sum(new_alpha[..., None] * nn_XtoP, dim=-2) / \
eps_denom(torch.sum(new_alpha, dim=-1, keepdim=True))
term_repul = repulsion_mu * torch.sum(new_beta[..., None] * delta, dim=-2) / \
eps_denom(torch.sum(new_beta, dim=-1, keepdim=True))
X = term_data + term_repul
if is_pointclouds(X0):
return X0.update_padded(X)
return X
def resample_uniformly(pointclouds,
neighborhood_size=8,
iters=1,
knn=None,
normals=None,
reproject=False,
repulsion_mu=1.0):
""" resample sample_iters times """
import math
import frnn
points_init, num_points = convert_pointclouds_to_tensor(pointclouds)
batch_size = num_points.shape[0]
# knn_result = knn_points(
# points_init, points_init, num_points, num_points, K=neighborhood_size + 1, return_nn=True)
diag = (points_init.view(-1, 3).max(dim=0).values -
points_init.view(-1, 3).min(0).values).norm().item()
avg_spacing = math.sqrt(diag / points_init.shape[1])
search_radius = min(4 * avg_spacing * neighborhood_size, 0.2)
if knn is None:
dists, idxs, _, grid = frnn.frnn_grid_points(points_init,
points_init,
num_points,
num_points,
K=neighborhood_size + 1,
r=search_radius,
grid=None,
return_nn=False)
knn = _KNN(dists=dists[..., 1:], idx=idxs[..., 1:], knn=None)
# estimate normals
if isinstance(pointclouds, torch.Tensor):
normals = normals
else:
normals = pointclouds.normals_padded()
if normals is None:
normals = estimate_pointcloud_normals(
points_init,
neighborhood_size=neighborhood_size,
disambiguate_directions=False)
else:
normals = F.normalize(normals, dim=-1)
points = points_init
for i in range(iters):
if reproject:
normals = denoise_normals(points,
normals,
num_points,
knn_result=knn)
points = project_to_latent_surface(points,
normals,
max_proj_iters=2,
max_est_iter=3)
if i > 0 and i % 3 == 0:
dists, idxs, _, grid = frnn.frnn_grid_points(points_init,
points_init,
num_points,
num_points,
K=neighborhood_size +
1,
r=search_radius,
grid=None,
return_nn=False)
knn = _KNN(dists=dists[..., 1:], idx=idxs[..., 1:], knn=None)
nn = frnn.frnn_gather(points, knn.idx, num_points)
pts_diff = points.unsqueeze(-2) - nn
dists = torch.sum(pts_diff**2, dim=-1)
knn_result = _KNN(dists=dists, idx=knn.idx, knn=nn)
deltap = knn_result.dists
inv_sigma_spatial = num_points / 2.0 / 16
spatial_w = torch.exp(-deltap * inv_sigma_spatial)
spatial_w[knn_result.idx < 0] = 0
# density_w = torch.sum(spatial_w, dim=-1) + 1.0
# 0.5 * derivative of (-r)exp(-r^2*inv)
density = frnn.frnn_gather(
spatial_w.sum(-1, keepdim=True) + 1.0, knn.idx, num_points)
nn_normals = frnn.frnn_gather(normals, knn_result.idx, num_points)
pts_diff_proj = pts_diff - (pts_diff * nn_normals).sum(
dim=-1, keepdim=True) * nn_normals
# move = 0.5 * torch.sum(density*spatial_w[..., None] * pts_diff_proj, dim=-2) / torch.sum(density.view_as(spatial_w)*spatial_w, dim=-1).unsqueeze(-1)
# move = F.normalize(move, dim=-1) * move.norm(dim=-1, keepdim=True).clamp_max(2*avg_spacing)
move = repulsion_mu * avg_spacing * torch.mean(
density * spatial_w[..., None] *
F.normalize(pts_diff_proj, dim=-1),
dim=-2)
points = points + move
# then project to latent surface again
if is_pointclouds(pointclouds):
return pointclouds.update_padded(points)
return points
import torch
import frnn
from pytorch_points.utils.pc_utils import read_ply
if __name__ == "__main__":
gpu = torch.device("cuda:0")
pc = torch.cuda.FloatTensor(read_ply("drill/drill_shaft_vrip.ply")[None, ...])
# print(pc.shape)
# print(pc.min(dim=1)[0], pc.max(dim=1)[0])
# pc -= pc.min(dim=1)[0]
# pc /= pc.max()
# print(pc.min(dim=1)[0], pc.max(dim=1)[0])
lengths = torch.ones((1,), dtype=torch.long, device=gpu) * pc.shape[1]
# print(lengths)
dists, idxs, nn, grid = frnn.frnn_grid_points(pc, pc, lengths, lengths, 8, 0.1, None, False, True, 2.0)
print(dists)
print(idxs)
def project_to_latent_surface(points,
normals,
sharpness_angle=60,
neighborhood_size=31,
max_proj_iters=10,
max_est_iter=5):
"""
RIMLS
"""
points, num_points = convert_pointclouds_to_tensor(points)
normals = F.normalize(normals, dim=-1)
sharpness_sigma = 1 - math.cos(sharpness_angle / 180 * math.pi)
diag = (points.max(dim=-2)[0] - points.min(dim=-2)[0]).norm(dim=-1)
avg_spacing = math.sqrt(diag / points.shape[1])
search_radius = min(16 * avg_spacing * neighborhood_size, 0.2)
dists, idxs, _, grid = frnn.frnn_grid_points(points,
points,
num_points,
num_points,
K=neighborhood_size + 1,
r=search_radius,
grid=None,
return_nn=False)
knn_result = _KNN(dists=dists[..., 1:], idx=idxs[..., 1:], knn=None)
# knn_normals = knn_gather(normals, knn_result.idx, num_points)
knn_normals = frnn.frnn_gather(normals, knn_result.idx, num_points)
inv_sigma_spatial = 1 / knn_result.dists[..., 0] / 16
# spatial_dist = 16 / inv_sigma_spatial
not_converged = torch.full(points.shape[:-1],
True,
device=points.device,
dtype=torch.bool)
itt = 0
it = 0
while True:
knn_pts = frnn.frnn_gather(points, knn_result.idx, num_points)
pts_diff = points[not_converged].unsqueeze(-2) - knn_pts[not_converged]
fx = torch.sum(pts_diff * knn_normals[not_converged], dim=-1)
not_converged_1 = torch.full(fx.shape[:-1],
True,
dtype=torch.bool,
device=fx.device)
knn_normals_1 = knn_normals[not_converged]
inv_sigma_spatial_1 = inv_sigma_spatial[not_converged]
f = points.new_zeros(points[not_converged].shape[:-1],
device=points.device)
grad_f = points.new_zeros(points[not_converged].shape,
device=points.device)
alpha = torch.ones_like(fx)
for itt in range(max_est_iter):
if itt > 0:
alpha_old = alpha
weights_n = (
(knn_normals_1[not_converged_1] -
grad_f[not_converged_1].unsqueeze(-2)).norm(dim=-1) /
0.5)**2
weights_n = torch.exp(-weights_n)
weights_p = torch.exp(-(
(fx[not_converged_1] -
f[not_converged_1].unsqueeze(-1))**2 *
inv_sigma_spatial_1[not_converged_1].unsqueeze(-1) / 4))
alpha[not_converged_1] = weights_n * weights_p
not_converged_1[not_converged_1] = (
alpha[not_converged_1] -
alpha_old[not_converged_1]).abs().max(dim=-1)[0] < 1e-4
if not not_converged_1.any():
break
deltap = torch.sum(pts_diff[not_converged_1] *
pts_diff[not_converged_1],
dim=-1)
phi = torch.exp(-deltap *
inv_sigma_spatial_1[not_converged_1].unsqueeze(-1))
# phi[deltap > spatial_dist] = 0
dphi = inv_sigma_spatial_1[not_converged_1].unsqueeze(-1) * phi
weights = phi * alpha[not_converged_1]
grad_weights = 2 * pts_diff * (dphi * weights).unsqueeze(-1)
sum_grad_weights = torch.sum(grad_weights, dim=-2)
sum_weight = torch.sum(weights, dim=-1)
sum_f = torch.sum(fx[not_converged_1] * weights, dim=-1)
sum_Gf = torch.sum(grad_weights *
fx[not_converged_1].unsqueeze(-1),
dim=-2)
sum_N = torch.sum(weights.unsqueeze(-1) *
knn_normals_1[not_converged_1],
dim=-2)
tmp_f = sum_f / eps_denom(sum_weight)
tmp_grad_f = (sum_Gf - tmp_f.unsqueeze(-1) * sum_grad_weights +
sum_N) / eps_denom(sum_weight).unsqueeze(-1)
grad_f[not_converged_1] = tmp_grad_f
f[not_converged_1] = tmp_f
move = f.unsqueeze(-1) * grad_f
points[not_converged] = points[not_converged] - move
mask = move.norm(dim=-1) > 5e-4
not_converged[not_converged] = mask
it = it + 1
if not not_converged.any() or it >= max_proj_iters:
break
return points
