def test_invalid_inputs_shapes(self, device="cuda:0"):
with self.assertRaisesRegex(
ValueError, "input can only be 2-dimensional."
):
values = torch.rand((100, 50, 2), device=device)
first_idxs = torch.tensor([0, 80], dtype=torch.int64, device=device)
packed_to_padded(values, first_idxs, 100)
with self.assertRaisesRegex(
ValueError, "input can only be 3-dimensional."
):
values = torch.rand((100,), device=device)
first_idxs = torch.tensor([0, 80], dtype=torch.int64, device=device)
padded_to_packed(values, first_idxs, 20)
with self.assertRaisesRegex(
ValueError, "input can only be 3-dimensional."
):
values = torch.rand((100, 50, 2, 2), device=device)
first_idxs = torch.tensor([0, 80], dtype=torch.int64, device=device)
padded_to_packed(values, first_idxs, 20)
def _test_padded_to_packed_helper(self, D, device):
"""
Check the results from packed_to_padded and PyTorch implementations
are the same.
"""
meshes = self.init_meshes(16, 100, 300, device=device)
mesh_to_faces_packed_first_idx = meshes.mesh_to_faces_packed_first_idx()
num_faces_per_mesh = meshes.num_faces_per_mesh()
max_faces = num_faces_per_mesh.max().item()
if D == 0:
values = torch.rand((len(meshes), max_faces), device=device)
else:
values = torch.rand((len(meshes), max_faces, D), device=device)
for i, num in enumerate(num_faces_per_mesh):
values[i, num:] = 0
values.requires_grad = True
values_torch = values.detach().clone()
values_torch.requires_grad = True
values_packed = padded_to_packed(
values,
mesh_to_faces_packed_first_idx,
num_faces_per_mesh.sum().item(),
)
values_packed_torch = TestPackedToPadded.padded_to_packed_python(
values_torch,
mesh_to_faces_packed_first_idx,
num_faces_per_mesh.sum().item(),
device,
)
# check forward
self.assertClose(values_packed, values_packed_torch)
# check backward
if D == 0:
grad_inputs = torch.rand(
(num_faces_per_mesh.sum().item()), device=device
)
else:
grad_inputs = torch.rand(
(num_faces_per_mesh.sum().item(), D), device=device
)
values_packed.backward(grad_inputs)
grad_outputs = values.grad
values_packed_torch.backward(grad_inputs)
grad_outputs_torch1 = values_torch.grad
grad_outputs_torch2 = TestPackedToPadded.packed_to_padded_python(
grad_inputs,
mesh_to_faces_packed_first_idx,
values.size(1),
device=device,
)
self.assertClose(grad_outputs, grad_outputs_torch1)
self.assertClose(grad_outputs, grad_outputs_torch2)
def compute(self, pointclouds: PointClouds3D, neighborhood_size=None):
if neighborhood_size is None:
neighborhood_size = self.neighborhood_size
num_points = pointclouds.num_points_per_cloud()
normals_packed = pointclouds.normals_packed()
assert (normals_packed is not None)
normals_ref = estimate_pointcloud_normals(
pointclouds, neighborhood_size=neighborhood_size)
normals_ref = padded_to_packed(normals_ref,
pointclouds.cloud_to_packed_first_idx(),
num_points.sum().item())
return 1 - F.cosine_similarity(normals_packed, normals_ref).abs()
def _compute_anisotropic_Vrk(self, pointclouds, **kwargs):
"""
determine the variance in the local surface frame h * Sk.T @ Sk based on curvature.
Args:
points_packed (Pointclouds3D): point clouds in object coordinates
Returns:
Vr (N, 3, 3): V_k^r matrix packed
"""
with torch.autograd.enable_grad():
points_padded = pointclouds.points_padded()
num_points = pointclouds.num_points_per_cloud()
# eigenvectors (=principal directions) in an ascending order of their
# corresponding eigenvalues, while the smallest eigenvalue's eigenvector
# corresponds to the normal direction
curvatures, local_frame = estimate_pointcloud_local_coord_frames(
pointclouds, neighborhood_size=8, disambiguate_directions=False)
local_frame = ops3d.padded_to_packed(
local_frame.reshape(local_frame.shape[:2] + (-1, )),
pointclouds.cloud_to_packed_first_idx(),
num_points.sum().item())
curvatures = ops3d.padded_to_packed(
curvatures, pointclouds.cloud_to_packed_first_idx(),
num_points.sum().item())
local_frame = local_frame.view(-1, 3, 3)[:, :, 1:]
# curvature only determines ratio of the two principle axis
# the actual size is based on a global max_size
curvatures = curvatures.view(-1, 3)[:, 1:]
curvature_ratios = curvatures / curvatures[:, -1:]
# TODO: compute density
curvatures = curvatures
Vr = local_frame @ torch.diag_embed(
curvatures) @ local_frame.transpose(1, 2)
return Vr, local_frame.transpose(1, 2)
def _denoise_normals(self,
point_clouds,
weights,
point_clouds_filter=None):
"""
robust normal mollification (Sec 4.4), i.e. replace normals with a weighted average
from neighboring normals
do this only for invisible points (?)
Args:
weights (tensors): (N,max_P,K)
"""
lengths = point_clouds.num_points_per_cloud()
P_total = lengths.sum().item()
normals = point_clouds.normals_padded()
batch_size, max_P, _ = normals.shape
knn_normals = ops.knn_gather(normals, self.knn_tree.idx, lengths)
normals_denoised = torch.sum(knn_normals * weights[:, :, :, None], dim=-2) / \
eps_denom(torch.sum(weights, dim=-1, keepdim=True))
# get point visibility so that we update only the non-visible or out-of-mask normals
if point_clouds_filter is not None:
try:
reliable_normals_mask = point_clouds_filter.visibility & point_clouds_filter.inmask
if len(point_clouds) != reliable_normals_mask.shape[0]:
if len(point_clouds
) == 1 and reliable_normals_mask.shape[0] > 1:
reliable_normals_mask = reliable_normals_mask.any(
dim=0, keepdim=True)
else:
ValueError(
"Incompatible point clouds {} and mask {}".format(
len(point_clouds),
reliable_normals_mask.shape))
# found visibility 0/1 as the last dimension of the features
# reset visible points normals to its original ones
normals_denoised[pts_reliable_normals_mask == 1] = normals[
reliable_normals_mask == 1]
except KeyError as e:
pass
normals_packed = point_clouds.normals_packed()
normals_denoised_packed = ops.padded_to_packed(
normals_denoised, point_clouds.cloud_to_packed_first_idx(),
P_total)
point_clouds.update_normals_(normals_denoised_packed)
return point_clouds
def _filter_points_with_invalid_depth(self, point_clouds, **kwargs):
self.cameras = kwargs.get('cameras', self.cameras)
lengths = point_clouds.num_points_per_cloud()
points_padded = point_clouds.points_padded()
with torch.autograd.no_grad():
to_view = self.cameras.get_world_to_view_transform()
points = to_view.transform_points(points_padded)
znear = getattr(self.cameras, 'znear', kwargs.get('znear', 1.0))
zfar = getattr(self.cameras, 'zfar', kwargs.get('zfar', 100.0))
mask = (points[..., 2] >= znear) & (points[..., 2] <= zfar)
mask_packed = ops3d.padded_to_packed(
mask.float(), point_clouds.cloud_to_packed_first_idx(),
lengths.sum().item()).bool()
if torch.all(mask_packed):
return point_clouds, mask_packed
# create point clouds again from packed
points_padded = point_clouds.points_padded()
normals_padded = point_clouds.normals_padded()
features_padded = point_clouds.features_padded()
# tuple to list, since pointclouds class doesn't accept tuples
points_list = [
points_padded[b][mask[b]] for b in range(points_padded.shape[0])
]
normals_list = features_list = None
if normals_padded is not None:
normals_list = [
normals_padded[b][mask[b]]
for b in range(normals_padded.shape[0])
]
if features_padded is not None:
features_list = [
features_padded[b][mask[b]]
for b in range(features_padded.shape[0])
]
new_point_clouds = point_clouds.__class__(points=points_list,
normals=normals_list,
features=features_list)
# for k in per_point_info:
# per_point_info[k] = per_point_info[k][mask_packed]
return new_point_clouds, mask_packed
def _filter_backface_points(self, point_clouds, **kwargs):
with torch.autograd.no_grad():
normals_view = self.transform_normals(point_clouds, **kwargs)
# mask = (normals_view[:, :, 2] < 1e-3) # error buffer
mask = normals_view[:, :, 2] < 0
lengths = point_clouds.num_points_per_cloud()
mask_packed = ops3d.padded_to_packed(
mask.float(), point_clouds.cloud_to_packed_first_idx(),
lengths.sum().item()).bool()
if torch.all(mask_packed):
return point_clouds, mask_packed
# create point clouds again from packed
points_padded = point_clouds.points_padded()
normals_padded = point_clouds.normals_padded()
features_padded = point_clouds.features_padded()
# tuple to list, since pointclouds class doesn't accept tuples
points_list = [
points_padded[b][mask[b]] for b in range(points_padded.shape[0])
]
normals_list = [
normals_padded[b][mask[b]] for b in range(normals_padded.shape[0])
]
if features_padded is not None:
features_list = [
features_padded[b][mask[b]]
for b in range(features_padded.shape[0])
]
new_point_clouds = point_clouds.__class__(points=points_list,
normals=normals_list,
features=features_list)
else:
new_point_clouds = point_clouds.__class__(points=points_list,
normals=normals_list)
# for k in per_point_info:
# per_point_info[k] = per_point_info[k][mask_packed]
return new_point_clouds, mask_packed
def backward(ctx, idx_grad, zbuf_grad, qvalue_grad, occ_grad):
# idx_grad and zbuf_grad are None (unless maybe we make weights depend on z? i.e. volumetric splatting)
grad_radii = None
grad_cloud_to_packed_first_idx = None
grad_num_points_per_cloud = None
grad_cutoff_thres = None
grad_depth_merging_thres = None
grad_image_size = None
grad_points_per_pixel = None
grad_bin_size = None
grad_max_points_per_bin = None
grad_radii_s = None
grad_backward_rbf = None
grads = (grad_cutoff_thres, grad_radii, grad_cloud_to_packed_first_idx,
grad_num_points_per_cloud, grad_depth_merging_thres,
grad_image_size, grad_points_per_pixel, grad_bin_size,
grad_max_points_per_bin, grad_radii_s, grad_backward_rbf)
radii_s = ctx.radii_backward_scaler
# either use OccRBFBackward or use OccBackward
pts_screen, ellipse_param, cutoff_threshold, radii, idx, zbuf0, \
cloud_to_packed_first_idx, num_points_per_cloud, \
= ctx.saved_tensors
depth_merging_threshold = ctx.depth_merging_threshold
backward_occ_fast = True
if not backward_occ_fast:
device = pts_screen.device
grads_input_xy = pts_screen.new_zeros((pts_screen.shape[0], 2))
grads_input_z = pts_screen.new_zeros((pts_screen.shape[0], 1))
mask = (idx[..., 0] >= 0).bool() # float
pts_visibility = torch.full((pts_screen.shape[0], ),
False,
dtype=torch.bool,
device=pts_screen.device)
# all rendered points (indices in packed points)
visible_idx = idx[mask].unique().long().view(-1)
visible_idx = visible_idx[visible_idx >= 0]
pts_visibility[visible_idx] = True
num_points_per_cloud = torch.stack([
x.sum() for x in torch.split(
pts_visibility, num_points_per_cloud.tolist(), dim=0)
])
cloud_to_packed_first_idx = num_points_2_cloud_to_packed_first_idx(
num_points_per_cloud)
pts_screen = pts_screen[pts_visibility]
radii = radii[pts_visibility]
grad_visible = _C._splat_points_occ_backward(
pts_screen, radii, occ_grad, cloud_to_packed_first_idx,
num_points_per_cloud, radii_s, depth_merging_threshold)
if torch.isnan(grad_visible).any(
) or not torch.isfinite(grad_visible).all():
print('invalid grad_visible')
assert (pts_visibility.sum() == grad_visible.shape[0])
grads_input_xy[pts_visibility] = grad_visible
_C._backward_zbuf(idx, zbuf_grad, grads_input_z)
# TODO necessary to concatenate
grads_input = torch.cat([grads_input_xy, grads_input_z], dim=-1)
else:
"""
We only care about rasterized points (visible points)
1. Filter [P,*] data to [P_visible,*] data
2. Fast backward cuda
2a. call FRNN insertion
2b. count_sort
"""
device = pts_screen.device
mask = (idx[..., 0] >= 0).bool() # float
pts_visibility = torch.full((pts_screen.shape[0], ),
False,
dtype=torch.bool,
device=pts_screen.device)
# all rendered points (indices in packed points)
visible_idx = idx[mask].unique().long().view(-1)
visible_idx = visible_idx[visible_idx >= 0]
pts_visibility[visible_idx] = True
num_points_per_cloud = torch.stack([
x.sum() for x in torch.split(
pts_visibility, num_points_per_cloud.tolist(), dim=0)
])
cloud_to_packed_first_idx = num_points_2_cloud_to_packed_first_idx(
num_points_per_cloud)
pts_screen_visible = pts_screen[pts_visibility]
radii_visible = radii[pts_visibility]
#####################################
# 2a. call FRNN insertion
#####################################
N = num_points_per_cloud.shape[0]
P = pts_screen_visible.shape[0]
assert (num_points_per_cloud.sum().item() == P)
# from frnn.frnn import GRID_PARAMS_SIZE, MAX_RES, prefix_sum_cuda
# imported from
from prefix_sum import prefix_sum_cuda
GRID_2D_PARAMS_SIZE = 6
GRID_2D_MAX_RES = 1024
GRID_2D_DELTA = 2
GRID_2D_TOTAL = 5
RADIUS_CELL_RATIO = 2
# first convert to padded
max_P = num_points_per_cloud.max().item()
pts_padded = ops3d.packed_to_padded(pts_screen_visible,
cloud_to_packed_first_idx,
max_P)
radii_padded = ops3d.packed_to_padded(radii_visible,
cloud_to_packed_first_idx,
max_P)
# determine search radius as max(radii)*radii_s
search_radius = torch.tensor([
radii_padded[i, :num_points_per_cloud[i]].median() * radii_s
for i in range(N)
],
dtype=torch.float,
device=device)
# create grid from scratch
# setup grid params
grid_params_cuda = torch.zeros((N, GRID_2D_PARAMS_SIZE),
dtype=torch.float,
device=pts_padded.device)
G = -1
pts_padded_2D = pts_padded[:, :, :2].clone().contiguous()
for i in range(N):
# 0-2 grid_min; 3 grid_delta; 4-6 grid_res; 7 grid_total
grid_min = pts_padded_2D[i, :num_points_per_cloud[i]].min(
dim=0)[0]
grid_max = pts_padded_2D[i, :num_points_per_cloud[i]].max(
dim=0)[0]
grid_params_cuda[i, :GRID_2D_DELTA] = grid_min
grid_size = grid_max - grid_min
cell_size = search_radius[i].item() / RADIUS_CELL_RATIO
if cell_size < grid_size.min() / GRID_2D_MAX_RES:
cell_size = grid_size.min() / GRID_2D_MAX_RES
grid_params_cuda[i, GRID_2D_DELTA] = 1 / cell_size
grid_params_cuda[i, GRID_2D_DELTA +
1:GRID_2D_TOTAL] = torch.floor(
grid_size / cell_size) + 1
grid_params_cuda[i, GRID_2D_TOTAL] = torch.prod(
grid_params_cuda[i, GRID_2D_DELTA + 1:GRID_2D_TOTAL])
if G < grid_params_cuda[i, GRID_2D_TOTAL]:
G = int(grid_params_cuda[i, GRID_2D_TOTAL].item())
# insert points into the grid
pc_grid_cnt = torch.zeros((N, G), dtype=torch.int, device=device)
pc_grid_cell = torch.full((N, max_P),
-1,
dtype=torch.int,
device=device)
pc_grid_idx = torch.full((N, max_P),
-1,
dtype=torch.int,
device=device)
frnn._C.insert_points_cuda(pts_padded_2D, num_points_per_cloud,
grid_params_cuda, pc_grid_cnt,
pc_grid_cell, pc_grid_idx, G)
# use prefix_sum from Matt Dean
grid_params = grid_params_cuda.cpu()
pc_grid_off = torch.full((N, G), 0, dtype=torch.int, device=device)
for i in range(N):
prefix_sum_cuda(pc_grid_cnt[i], grid_params[i, GRID_2D_TOTAL],
pc_grid_off[i])
# sort points according to their grid positions and insertion orders
# sort based on x, y first. Then we will use points_sorted_idxs to recover the points_sorted with Z
points_sorted = torch.zeros((N, max_P, 2),
dtype=torch.float,
device=device)
points_sorted_idxs = torch.full((N, max_P),
-1,
dtype=torch.int,
device=device)
frnn._C.counting_sort_cuda(
pts_padded_2D,
num_points_per_cloud,
pc_grid_cell,
pc_grid_idx,
pc_grid_off,
points_sorted, # (N,P,2)
points_sorted_idxs # (N,P)
)
new_points_sorted = torch.zeros_like(pts_padded)
for i in range(N):
points_sorted_idxs_i = points_sorted_idxs[
i, :num_points_per_cloud[i]].long().unsqueeze(1).expand(
-1, 3)
new_points_sorted[i, :num_points_per_cloud[i]] = torch.gather(
pts_padded[i], 0, points_sorted_idxs_i)
# print(points_sorted[i, :10])
# print(new_points_sorted[i, :10])
# new_points_sorted = torch.gather(pts_padded, 1, points_sorted_idxs.long().unsqueeze(2).expand(-1, -1, 3))
assert (new_points_sorted is not None and pc_grid_off is not None
and points_sorted_idxs is not None
and grid_params_cuda is not None)
# convert sorted_points and sorted_points_idxs to packed (P, )
points_sorted = ops3d.padded_to_packed(new_points_sorted,
cloud_to_packed_first_idx,
P)
# padded_to_packed only supports torch.float32...
shifted_points_sorted_idxs = points_sorted_idxs + cloud_to_packed_first_idx.float(
).unsqueeze(1)
points_sorted_idxs = ops3d.padded_to_packed(
shifted_points_sorted_idxs, cloud_to_packed_first_idx, P)
points_sorted_idxs_2D = points_sorted_idxs.long().unsqueeze(
1).expand(-1, 2)
radii_sorted = torch.gather(radii_visible, 0,
points_sorted_idxs_2D)
pc_grid_off += cloud_to_packed_first_idx.unsqueeze(1)
grad_sorted = _C._splat_points_occ_fast_cuda_backward(
points_sorted, radii_sorted, search_radius, occ_grad,
num_points_per_cloud, cloud_to_packed_first_idx, pc_grid_off,
grid_params_cuda)
# grad_sorted_slow = _C._splat_points_occ_backward(points_sorted, radii_sorted,
# occ_grad, cloud_to_packed_first_idx, num_points_per_cloud,
# radii_s, depth_merging_threshold)
# breakpoint()
# points_sorted_idxs_3D = points_sorted_idxs.long().unsqueeze(1).expand(-1, 3)
# print(points_sorted_idxs_3D.max(), grad_sorted.shape[0])
grad_visible = torch.zeros_like(grad_sorted).scatter_(
0, points_sorted_idxs_2D, grad_sorted)
# grad_visible_slow = _C._splat_points_occ_backward(pts_screen[pts_visibility], radii[pts_visibility],
# occ_grad, cloud_to_packed_first_idx, num_points_per_cloud,
# radii_s, depth_merging_threshold)
# breakpoint()
if torch.isnan(grad_visible).any(
) or not torch.isfinite(grad_visible).all():
print('invalid grad_visible')
assert (pts_visibility.sum() == grad_visible.shape[0])
grads_input_xy = pts_screen.new_zeros(pts_screen.shape[0], 2)
grads_input_z = pts_screen.new_zeros(pts_screen.shape[0], 1)
# print("1")
grads_input_xy[pts_visibility] = grad_visible
_C._backward_zbuf(idx, zbuf_grad, grads_input_z)
grads_input = torch.cat([grads_input_xy, grads_input_z], dim=-1)
# print("2")
pts_grad = grads_input
return (pts_grad, None) + grads
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 compute(self,
point_clouds: PointClouds3D,
points_filters=None,
rebuild_knn=True,
**kwargs):
self.knn_tree = kwargs.get('knn_tree', self.knn_tree)
self.knn_mask = kwargs.get('knn_mask', self.knn_mask)
lengths = point_clouds.num_points_per_cloud()
P_total = lengths.sum().item()
points_padded = point_clouds.points_padded()
# Compute necessary weights to project points to local plane
# TODO(yifan): This part is same as ProjectionLoss
# how can we at best save repetitive computation
with torch.autograd.no_grad():
if rebuild_knn or self.knn_tree is None or points_padded.shape[:
2] != self.knn_tree.shape[:
2]:
self._build_knn(point_clouds)
phi = self.get_phi(point_clouds, **kwargs)
self._denoise_normals(point_clouds, phi, points_filters)
# compute wn and wr
# TODO(yifan): visibility weight?
normal_w = self.get_normal_w(point_clouds, **kwargs)
# update normals for a second iteration (?) Eq.(10)
point_clouds = self._denoise_normals(point_clouds, phi * normal_w,
points_filters)
# compose weights
weights = phi * normal_w
weights[~self.knn_mask] = 0
# outside filter_scale*local_point_spacing weights
mask_ball_query = self.knn_tree.dists > (
self.filter_scale * self.knn_tree.dists[:, :, :1] * 2.0)
weights[mask_ball_query] = 0.0
# project the point to a local surface
knn_normals = ops.knn_gather(point_clouds.normals_padded(),
self.knn_tree.idx, lengths)
dist_to_surface = torch.sum(
(self.knn_tree.knn.detach() - points_padded.unsqueeze(-2)) *
knn_normals,
dim=-1)
deltap = torch.sum(
dist_to_surface[..., None] * weights[..., None] * knn_normals,
dim=-2) / eps_denom(torch.sum(weights, dim=-1, keepdim=True))
points_projected = points_padded + deltap
if get_debugging_mode():
# points_padded.requires_grad_(True)
def save_grad():
lengths = point_clouds.num_points_per_cloud()
def _save_grad(grad):
dbg_tensor = get_debugging_tensor()
if dbg_tensor is None:
logger_py.error("dbg_tensor is None")
if grad is None:
logger_py.error('grad is None')
# a dict of list of tensors
dbg_tensor.pts_world_grad['repel'] = [
grad[b, :lengths[b]].detach().cpu()
for b in range(grad.shape[0])
]
return _save_grad
dbg_tensor = get_debugging_tensor()
dbg_tensor.pts_world['repel'] = [
points_padded[b, :lengths[b]].detach().cpu()
for b in range(points_padded.shape[0])
]
handle = points_padded.register_hook(save_grad())
self.hooks.append(handle)
with torch.autograd.no_grad():
spatial_w = self.get_spatial_w(point_clouds, points_projected)
# density_w = self.get_density_w(point_clouds)
# density weight is actually spatial_w + 1
density_w = torch.sum(spatial_w, dim=-1, keepdim=True) + 1.0
weights = normal_w * spatial_w * density_w
weights[~self.knn_mask] = 0
weights[mask_ball_query] = 0
deltap = points_projected[:, :, None, :] - self.knn_tree.knn.detach()
point_to_point_dist = torch.sum(deltap * deltap, dim=-1)
# convert everything to packed
weights = ops.padded_to_packed(
weights, point_clouds.cloud_to_packed_first_idx(), P_total)
point_to_point_dist = ops.padded_to_packed(
point_to_point_dist, point_clouds.cloud_to_packed_first_idx(),
P_total)
# we want to maximize this, so negative sign
point_to_point_dist = -torch.sum(point_to_point_dist * weights,
dim=1) / eps_denom(
torch.sum(weights, dim=1))
return point_to_point_dist
def compute(self,
point_clouds: PointClouds3D,
points_filters=None,
rebuild_knn=False,
**kwargs):
"""
Args:
point_clouds
(optional) knn_tree: output from ops.knn_points excluding the query point itself
(optional) knn_mask: mask valid knn results
Returns:
(P, N)
"""
self.sharpness_sigma = kwargs.get('sharpness_sigma',
self.sharpness_sigma)
self.filter_scale = kwargs.get('filter_scale', self.filter_scale)
self.knn_tree = kwargs.get('knn_tree', self.knn_tree)
self.knn_mask = kwargs.get('knn_mask', self.knn_mask)
lengths = point_clouds.num_points_per_cloud()
P_total = lengths.sum().item()
points = point_clouds.points_padded()
# - determine phi spatial with using local point spacing (i.e. 2*dist_to_nn)
# - denoise normals
# - determine w_normal
# - mask out values outside ballneighbor i.e. d > filterSpatialScale * localPointSpacing
# - projected distance dot(ni, x-xi)
# - multiply and normalize the weights
with torch.autograd.no_grad():
if rebuild_knn or self.knn_tree is None or self.knn_tree.idx.shape[:
2] != points.shape[:
2]:
self._build_knn(point_clouds)
phi = self.get_phi(point_clouds, **kwargs)
# robust normal mollification (Sec 4.4), i.e. replace normals with a weighted average
# from neighboring normals Eq.(11)
point_clouds = self._denoise_normals(point_clouds, phi,
points_filters)
# compute wn and wr
# TODO(yifan): visibility weight?
normal_w = self.get_normal_w(point_clouds, **kwargs)
spatial_w = self.get_spatial_w(point_clouds, **kwargs)
# update normals for a second iteration (?) Eq.(10)
point_clouds = self._denoise_normals(point_clouds, phi * normal_w,
points_filters)
# compose weights
weights = phi * spatial_w * normal_w
weights[~self.knn_mask] = 0
# outside filter_scale*local_point_spacing weights
mask_ball_query = self.knn_tree.dists > (
self.filter_scale * self.knn_tree.dists[:, :, :1] * 2.0)
weights[mask_ball_query] = 0.0
# (B, P, k), dot product distance to surface
# (we need to gather again because the normals have been changed in the denoising step)
knn_normals = ops.knn_gather(point_clouds.normals_padded(),
self.knn_tree.idx, lengths)
# if points.requires_grad:
# from DSS.core.rasterizer import _dbg_tensor
# def save_grad(name):
# def _save_grad(grad):
# _dbg_tensor[name] = grad.detach().cpu()
# return _save_grad
# points.register_hook(save_grad('proj_grad'))
dist_to_surface = torch.sum(
(self.knn_tree.knn.detach() - points.unsqueeze(-2)) * knn_normals,
dim=-1)
if get_debugging_mode():
# points.requires_grad_(True)
def save_grad():
lengths = point_clouds.num_points_per_cloud()
def _save_grad(grad):
dbg_tensor = get_debugging_tensor()
if dbg_tensor is None:
logger_py.error("dbg_tensor is None")
if grad is None:
logger_py.error('grad is None')
# a dict of list of tensors
dbg_tensor.pts_world_grad['proj'] = [
grad[b, :lengths[b]].detach().cpu()
for b in range(grad.shape[0])
]
return _save_grad
dbg_tensor = get_debugging_tensor()
dbg_tensor.pts_world['proj'] = [
points[b, :lengths[b]].detach().cpu()
for b in range(points.shape[0])
]
handle = points.register_hook(save_grad())
self.hooks.append(handle)
# convert everything to packed
weights = ops.padded_to_packed(
weights, point_clouds.cloud_to_packed_first_idx(), P_total)
dist_to_surface = ops.padded_to_packed(
dist_to_surface, point_clouds.cloud_to_packed_first_idx(), P_total)
# compute weighted signed distance to surface
dist_to_surface = torch.sum(weights * dist_to_surface,
dim=-1) / eps_denom(
torch.sum(weights, dim=-1))
loss = dist_to_surface * dist_to_surface
return loss
def compute(self,
point_clouds: PointClouds3D,
points_filter=None,
rebuild_knn=True,
**kwargs):
self.knn_tree = kwargs.get('knn_tree', self.knn_tree)
self.knn_mask = kwargs.get('knn_mask', self.knn_mask)
lengths = point_clouds.num_points_per_cloud()
P_total = lengths.sum().item()
points_padded = point_clouds.points_padded()
if not points_padded.requires_grad:
logger_py.warn(
'Computing repulsion loss, but points_padded is not differentiable.'
)
# Compute necessary weights to project points to local plane
# TODO(yifan): This part is same as ProjectionLoss
# how can we at best save repetitive computation
with torch.autograd.no_grad():
if rebuild_knn or self.knn_tree is None or points_padded.shape[:
2] != self.knn_tree.shape[:
2]:
self._build_knn(point_clouds)
phi = self.get_phi(point_clouds, **kwargs)
point_clouds = self._denoise_normals(point_clouds,
phi,
points_filter,
inplace=False)
# project the point to a local surface
knn_diff = points_padded.unsqueeze(-2) - self.knn_tree.knn.detach()
knn_normals = ops.knn_gather(point_clouds.normals_padded(),
self.knn_tree.idx, lengths)
pts_diff_proj = knn_diff - \
(knn_diff * knn_normals).sum(dim=-1, keepdim=True) * knn_normals
if get_debugging_mode():
# points_padded.requires_grad_(True)
def save_grad():
lengths = point_clouds.num_points_per_cloud()
def _save_grad(grad):
dbg_tensor = get_debugging_tensor()
if dbg_tensor is None:
logger_py.error("dbg_tensor is None")
if grad is None:
logger_py.error('grad is None')
# a dict of list of tensors
dbg_tensor.pts_world_grad['repel'] = [
grad[b, :lengths[b]].detach().cpu()
for b in range(grad.shape[0])
]
return _save_grad
if points_padded.requires_grad:
dbg_tensor = get_debugging_tensor()
dbg_tensor.pts_world['repel'] = [
points_padded[b, :lengths[b]].detach().cpu()
for b in range(points_padded.shape[0])
]
handle = points_padded.register_hook(save_grad())
self.hooks.append(handle)
with torch.autograd.no_grad():
spatial_w = self.get_spatial_w(point_clouds, **kwargs)
# set far neighbors' spatial_w to 0
normal_w = self.get_normal_w(point_clouds, **kwargs)
density_w = torch.sum(spatial_w, dim=-1, keepdim=True) + 1.0
weights = spatial_w * normal_w
# convert everything to packed
weights = ops.padded_to_packed(
weights, point_clouds.cloud_to_packed_first_idx(), P_total)
pts_diff_proj = ops.padded_to_packed(
pts_diff_proj.contiguous().view(pts_diff_proj.shape[0],
pts_diff_proj.shape[1], -1),
point_clouds.cloud_to_packed_first_idx(),
P_total).view(P_total, -1, 3)
density_w = ops.padded_to_packed(
density_w, point_clouds.cloud_to_packed_first_idx(), P_total)
# we want to maximize this, so negative sign
repel_vec = torch.sum(pts_diff_proj * weights.unsqueeze(-1),
dim=1) / eps_denom(
torch.sum(weights, dim=1).unsqueeze(-1))
repel_vec = repel_vec * density_w
loss = torch.exp(-repel_vec.abs())
# if get_debugging_mode():
# # save to dbg folder as normal
# from ..utils.io import save_ply
# save_ply('./dbg_repel_diff.ply', point_clouds.points_packed().cpu().detach(), normals=repel_vec.cpu().detach())
return loss
def compute(self,
point_clouds: PointClouds3D,
points_filter=None,
rebuild_knn=False,
**kwargs):
"""
Args:
point_clouds
(optional) knn_tree: output from ops.knn_points excluding the query point itself
(optional) knn_mask: mask valid knn results
Returns:
(P, N)
"""
self.sharpness_sigma = kwargs.get('sharpness_sigma',
self.sharpness_sigma)
self.filter_scale = kwargs.get('filter_scale', self.filter_scale)
self.knn_tree = kwargs.get('knn_tree', self.knn_tree)
self.knn_mask = kwargs.get('knn_mask', self.knn_mask)
lengths = point_clouds.num_points_per_cloud()
P_total = lengths.sum().item()
points = point_clouds.points_padded()
# - determine phi spatial with using local point spacing (i.e. 2*dist_to_nn)
# - denoise normals
# - determine w_normal
# - mask out values outside ballneighbor i.e. d > filterSpatialScale * localPointSpacing
# - projected distance dot(ni, x-xi)
# - multiply and normalize the weights
with torch.autograd.no_grad():
if rebuild_knn or self.knn_tree is None or self.knn_tree.idx.shape[:
2] != points.shape[:
2]:
self._build_knn(point_clouds)
phi = self.get_phi(point_clouds, **kwargs)
# robust normal mollification (Sec 4.4), i.e. replace normals with a weighted average
# from neighboring normals Eq.(11)
point_clouds = self._denoise_normals(point_clouds,
phi,
points_filter,
inplace=False)
# compute wn and wr
normal_w = self.get_normal_w(point_clouds, **kwargs)
# visibility weight
visibility_nb = ops.knn_gather(
points_filter.visibility.unsqueeze(-1), self.knn_tree.idx,
lengths)
visibility_w = visibility_nb.float()
visibility_w[~visibility_nb] = 0.1
# compose weights
weights = phi * normal_w * visibility_w.squeeze(-1)
# (B, P, k), dot product distance to surface
knn_normals = ops.knn_gather(point_clouds.normals_padded(),
self.knn_tree.idx, lengths)
if get_debugging_mode():
# points.requires_grad_(True)
def save_grad():
lengths = point_clouds.num_points_per_cloud()
def _save_grad(grad):
dbg_tensor = get_debugging_tensor()
if dbg_tensor is None:
logger_py.error("dbg_tensor is None")
if grad is None:
logger_py.error('grad is None')
# a dict of list of tensors
dbg_tensor.pts_world_grad['proj'] = [
grad[b, :lengths[b]].detach().cpu()
for b in range(grad.shape[0])
]
return _save_grad
if points.requires_grad:
dbg_tensor = get_debugging_tensor()
dbg_tensor.pts_world['proj'] = [
points[b, :lengths[b]].detach().cpu()
for b in range(points.shape[0])
]
handle = points.register_hook(save_grad())
self.hooks.append(handle)
sdf = torch.sum(
(self.knn_tree.knn.detach() - points.unsqueeze(-2)) * knn_normals,
dim=-1)
# convert everything to packed
weights = ops.padded_to_packed(
weights, point_clouds.cloud_to_packed_first_idx(), P_total)
sdf = ops.padded_to_packed(sdf,
point_clouds.cloud_to_packed_first_idx(),
P_total)
# if get_debugging_mode():
# # save to dbg folder as normal
# from ..utils.io import save_ply
# save_ply('./dbg_repel_diff.ply', point_clouds.points_packed().cpu().detach(), normals=repel_vec.cpu().detach())
distance_to_face = sdf * sdf
# compute weighted signed distance to surface
loss = torch.sum(weights * distance_to_face, dim=-1) / eps_denom(
torch.sum(weights, dim=-1))
return loss
def test_dataset(self):
# 1. rerender input point clouds / meshes using the saved camera_mat
# compare mask image with saved mask image
# 2. backproject masked points to space with dense depth map,
# fuse all views and save
batch_size = 1
device = torch.device('cuda:0')
data_dir = 'data/synthetic/cube_mesh'
output_dir = os.path.join('tests', 'outputs', 'test_data')
if not os.path.isdir(output_dir):
os.makedirs(output_dir)
# dataset
dataset = MVRDataset(data_dir=data_dir,
load_dense_depth=True,
mode="train")
data_loader = torch.utils.data.DataLoader(dataset,
batch_size=batch_size,
num_workers=0,
shuffle=False)
meshes = load_objs_as_meshes([os.path.join(data_dir,
'mesh.obj')]).to(device)
cams = dataset.get_cameras().to(device)
image_size = imageio.imread(dataset.image_files[0]).shape[0]
# initialize rasterizer, we check mask pngs only, so no need to create lights and shaders etc
raster_settings = RasterizationSettings(
image_size=image_size,
blur_radius=0.0,
faces_per_pixel=5,
bin_size=
None, # this setting controls whether naive or coarse-to-fine rasterization is used
max_faces_per_bin=None # this setting is for coarse rasterization
)
rasterizer = MeshRasterizer(cameras=None,
raster_settings=raster_settings)
# render with loaded cameras positions and training tranformation functions
pixel_world_all = []
for idx, data in enumerate(data_loader):
# get datas
img = data.get('img.rgb').to(device)
assert (img.min() >= 0 and img.max() <= 1
), "Image must be a floating number between 0 and 1."
mask_gt = data.get('img.mask').to(device).permute(0, 2, 3, 1)
camera_mat = data['camera_mat'].to(device)
cams.R, cams.T = decompose_to_R_and_t(camera_mat)
cams._N = cams.R.shape[0]
cams.to(device)
self.assertTrue(
torch.equal(cams.get_world_to_view_transform().get_matrix(),
camera_mat))
# transform to view and rerender with non-rotated camera
verts_padded = transform_to_camera_space(meshes.verts_padded(),
cams)
meshes_in_view = meshes.offset_verts(
-meshes.verts_packed() + padded_to_packed(
verts_padded, meshes.mesh_to_verts_packed_first_idx(),
meshes.verts_packed().shape[0]))
fragments = rasterizer(meshes_in_view,
cameras=dataset.get_cameras().to(device))
# compare mask
mask = fragments.pix_to_face[..., :1] >= 0
imageio.imwrite(os.path.join(output_dir, "mask_%06d.png" % idx),
mask[0, ...].cpu().to(dtype=torch.uint8) * 255)
# allow 5 pixels difference
self.assertTrue(torch.sum(mask_gt != mask) < 5)
# check dense maps
# backproject points to the world pixel range (-1, 1)
pixels = arange_pixels((image_size, image_size),
batch_size)[1].to(device)
depth_img = data.get('img.depth').to(device)
# get the depth and mask at the sampled pixel position
depth_gt = get_tensor_values(depth_img,
pixels,
squeeze_channel_dim=True)
mask_gt = get_tensor_values(mask.permute(0, 3, 1, 2).float(),
pixels,
squeeze_channel_dim=True).bool()
# get pixels and depth inside the masked area
pixels_packed = pixels[mask_gt]
depth_gt_packed = depth_gt[mask_gt]
first_idx = torch.zeros((pixels.shape[0], ),
device=device,
dtype=torch.long)
num_pts_in_mask = mask_gt.sum(dim=1)
first_idx[1:] = num_pts_in_mask.cumsum(dim=0)[:-1]
pixels_padded = packed_to_padded(pixels_packed, first_idx,
num_pts_in_mask.max().item())
depth_gt_padded = packed_to_padded(depth_gt_packed, first_idx,
num_pts_in_mask.max().item())
# backproject to world coordinates
# contains nan and infinite values due to depth_gt_padded containing 0.0
pixel_world_padded = transform_to_world(pixels_padded,
depth_gt_padded[..., None],
cams)
# transform back to list, containing no padded values
split_size = num_pts_in_mask[..., None].repeat(1, 2)
split_size[:, 1] = 3
pixel_world_list = padded_to_list(pixel_world_padded, split_size)
pixel_world_all.extend(pixel_world_list)
idx += 1
if idx >= 10:
break
pixel_world_all = torch.cat(pixel_world_all, dim=0)
mesh = trimesh.Trimesh(vertices=pixel_world_all.cpu(),
faces=None,
process=False)
mesh.export(os.path.join(output_dir, 'pixel_to_world.ply'))
