def test_generate_points(self, octrees, lengths, max_level):
out_level, pyramids, exsum = scan_octrees(octrees, lengths)
point_hierarchies = generate_points(octrees, pyramids, exsum)
expected_point_hierarchies = []
bits_t = uint8_to_bits(octrees).reshape(-1, 2, 2, 2).cpu()
octree_first_idx = 0
for bs, length in enumerate(lengths):
expected_point_hierarchies.append(
torch.tensor([[0, 0, 0]], dtype=torch.long))
cur_bits_t = bits_t[octree_first_idx:octree_first_idx + length]
offsets = torch.tensor([[0, 0, 0]], dtype=torch.int32)
for i in range(max_level):
next_offset = []
cur_level_num_nodes = pyramids[bs, 0, i]
level_first_idx = pyramids[bs, 1, i]
for cur_level_node_idx in range(cur_level_num_nodes):
node_bits = cur_bits_t[level_first_idx +
cur_level_node_idx]
offset = offsets[cur_level_node_idx]
point_coords = torch.nonzero(
node_bits, as_tuple=False) + offset.unsqueeze(0)
expected_point_hierarchies.append(point_coords)
next_offset.append(point_coords * 2)
offsets = torch.cat(next_offset, dim=0)
octree_first_idx += length
expected_point_hierarchies = torch.cat(expected_point_hierarchies,
dim=0).cuda().short()
assert torch.equal(point_hierarchies, expected_point_hierarchies)
def test_scan_octrees(self, octrees, lengths):
expected_pyramids = torch.tensor(
[[[1, 2, 3, 3, 0], [0, 1, 3, 6, 9]],
[[1, 1, 3, 13, 0], [0, 1, 2, 5, 18]]],
dtype=torch.int32)
expected_exsum = torch.tensor(
[0, 2, 4, 5, 6, 7, 8, 0, 1, 4, 5, 13, 17],
dtype=torch.int32,
device='cuda')
max_level, pyramids, exsum = scan_octrees(octrees, lengths)
assert max_level == 3
assert torch.equal(pyramids, expected_pyramids)
assert torch.equal(exsum, expected_exsum)
def test_feature_grids_to_spc(self, sparse_feature_grids,
expected_out_feature_grids, device):
octrees, lengths, features = feature_grids_to_spc(sparse_feature_grids)
assert octrees.device.type == device
assert features.device.type == device
octrees = octrees.cuda()
features = features.cuda()
max_level, pyramids, exsum = scan_octrees(octrees, lengths)
point_hierarchies = generate_points(octrees, pyramids, exsum)
out_feature_grids = to_dense(point_hierarchies, pyramids, features,
max_level)
assert torch.equal(out_feature_grids, expected_out_feature_grids)
def test_generate_points(self, octrees, lengths):
max_level, pyramids, exsum = scan_octrees(octrees, lengths)
expected_point_hierarchies = torch.tensor([
[0, 0, 0],
[0, 0, 0], [1, 0, 0],
[0, 0, 1], [0, 1, 0], [3, 0, 1],
[1, 1, 3], [1, 3, 1], [6, 1, 3],
[0, 0, 0],
[1, 1, 1],
[3, 2, 2], [3, 2, 3], [3, 3, 2],
[7, 4, 5], [6, 4, 6], [6, 4, 7], [6, 5, 6], [6, 5, 7], [7, 4, 6], \
[7, 4, 7], [7, 5, 6], [7, 5, 7], [6, 6, 4], [6, 7, 4], \
[7, 6, 4], [7, 7, 4]
], device='cuda', dtype=torch.int16)
point_hierarchies = generate_points(octrees, pyramids, exsum)
assert torch.equal(point_hierarchies, expected_point_hierarchies)
def test_ones(self, batch_size, feature_dim, height, width, depth, dtype,
device):
feature_grids = torch.ones(
(batch_size, feature_dim, height, width, depth),
dtype=dtype,
device=device)
octrees, lengths, features = feature_grids_to_spc(feature_grids)
assert octrees.device.type == device
assert features.device.type == device
octrees = octrees.cuda()
features = features.cuda()
max_level, pyramids, exsum = scan_octrees(octrees, lengths)
point_hierarchies = generate_points(octrees, pyramids, exsum)
out_feature_grids = to_dense(point_hierarchies, pyramids, features,
max_level)
assert torch.all(out_feature_grids[:, :, :height, :width, :depth] == 1)
assert torch.all(out_feature_grids[:, :, height:] == 0)
assert torch.all(out_feature_grids[:, :, :, width:] == 0)
assert torch.all(out_feature_grids[..., depth:] == 0)
def test_scan_octrees(self, octrees, lengths, max_level):
# Naive implementation
num_childrens_per_node = uint8_bits_sum(octrees).cpu()
octree_start_idx = 0
num_childrens_per_level = []
levels_first_idx = []
expected_exsum = torch.zeros(
(num_childrens_per_node.shape[0] + lengths.shape[0], ),
dtype=torch.int32)
for bs, length in enumerate(lengths):
cur_num_childrens_per_node = \
num_childrens_per_node[octree_start_idx:octree_start_idx + length]
num_childrens_per_level.append([1])
levels_first_idx.append([0])
for i in range(max_level):
cur_idx = levels_first_idx[-1][-1]
cur_num_childrens = num_childrens_per_level[-1][-1]
num_childrens_per_level[-1].append(
int(
torch.sum(
cur_num_childrens_per_node[cur_idx:cur_idx +
cur_num_childrens])))
levels_first_idx[-1].append(cur_idx + cur_num_childrens)
levels_first_idx[-1].append(levels_first_idx[-1][-1] +
num_childrens_per_level[-1][-1])
num_childrens_per_level[-1].append(0)
# + bs + 1 because torch.cumsum is inclusive
expected_exsum[octree_start_idx + bs + 1:octree_start_idx + bs + 1 + length] = \
torch.cumsum(cur_num_childrens_per_node, dim=0)
octree_start_idx += length
num_childrens_per_level = torch.tensor(num_childrens_per_level,
dtype=torch.int32)
levels_first_idx = torch.tensor(levels_first_idx, dtype=torch.int32)
expected_pyramids = torch.stack(
[num_childrens_per_level, levels_first_idx], dim=1)
expected_exsum = expected_exsum.cuda()
out_level, pyramids, exsum = scan_octrees(octrees, lengths)
assert out_level == max_level
assert torch.equal(pyramids, expected_pyramids)
assert torch.equal(exsum, expected_exsum)
def test_simple(self, with_spc_to_dict):
bits_t = torch.tensor(
[[0, 0, 0, 1, 0, 0, 0, 1], [0, 0, 0, 0, 0, 1, 1, 0],
[0, 0, 1, 0, 0, 0, 0, 0], [1, 0, 0, 0, 0, 0, 0, 0],
[1, 0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 1, 0, 0, 0],
[1, 0, 0, 0, 0, 0, 0, 0], [0, 1, 1, 1, 0, 0, 0, 0],
[0, 0, 1, 0, 0, 0, 0, 0], [1, 1, 1, 1, 1, 1, 1, 1],
[0, 1, 0, 1, 0, 1, 0, 1]],
device='cuda',
dtype=torch.float)
octrees = bits_to_uint8(torch.flip(bits_t, dims=(-1, )))
lengths = torch.tensor([6, 5], dtype=torch.int)
max_level, pyramids, exsum = scan_octrees(octrees, lengths)
point_hierarchies = generate_points(octrees, pyramids, exsum)
coalescent_features = torch.tensor([
1., 2., 3., 4., 5., 6., 7., 8., 9., 10., 11., 12., 13., 14., 15.,
16.
],
device='cuda',
dtype=torch.float).reshape(-1, 1)
feat_idx = torch.tensor(
[[0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1],
[1, 1, 6, 7, 6, 6, 6, 6, 7, 7, 7, 7, 6, 6, 7, 7],
[1, 3, 1, 4, 4, 4, 5, 5, 4, 4, 5, 5, 6, 7, 6, 7],
[3, 1, 3, 5, 6, 7, 6, 7, 6, 7, 6, 7, 4, 4, 4, 4]],
dtype=torch.long)
expected_feature_grids = torch.zeros((2, 1, 8, 8, 8),
dtype=torch.float,
device='cuda')
expected_feature_grids[feat_idx[0], :, feat_idx[1], feat_idx[2],
feat_idx[3]] = coalescent_features
if with_spc_to_dict:
feature_grids = to_dense(**Spc(octrees, lengths).to_dict(),
input=coalescent_features)
else:
feature_grids = to_dense(point_hierarchies, pyramids,
coalescent_features, max_level)
assert torch.equal(feature_grids, expected_feature_grids)
def test_to_dense(self, batch_size, max_level, feature_dim):
octrees, lengths = random_spc_octrees(batch_size, max_level, 'cuda')
max_level, pyramids, exsum = scan_octrees(octrees, lengths)
point_hierarchies = generate_points(octrees, pyramids, exsum)
in_num_nodes = torch.sum(pyramids[:, 0, -2])
coalescent_features = torch.rand((in_num_nodes, feature_dim),
device='cuda',
requires_grad=True)
expected_size = 2**max_level
feat_idx = []
bs_start_idx = 0
for bs in range(batch_size):
start_idx = pyramids[bs, 1, -2] + bs_start_idx
num_points = pyramids[bs, 0, -2]
feat_idx.append(
torch.nn.functional.pad(point_hierarchies[start_idx:start_idx +
num_points], (1, 0),
value=bs))
bs_start_idx += pyramids[bs, 1, -1]
feat_idx = torch.cat(feat_idx, dim=0).permute(1, 0).long()
expected_feature_grids = torch.zeros(
(batch_size, feature_dim, expected_size, expected_size,
expected_size),
device='cuda')
expected_feature_grids[feat_idx[0], :, feat_idx[1], feat_idx[2],
feat_idx[3]] = coalescent_features
# test forward
feature_grids = to_dense(point_hierarchies, pyramids,
coalescent_features, max_level)
assert torch.equal(expected_feature_grids, feature_grids)
grad_out = torch.rand_like(feature_grids)
feature_grids.backward(grad_out)
octrees, lengths, coalescent_expected_grad = feature_grids_to_spc(
grad_out, torch.any(feature_grids != 0, dim=1))
assert torch.equal(coalescent_features.grad, coalescent_expected_grad)
def test_query(self):
points = torch.tensor([[3, 2, 0], [3, 1, 1], [0, 0, 0], [3, 3, 3]],
device='cuda',
dtype=torch.short)
octree = unbatched_points_to_octree(points, 2)
length = torch.tensor([len(octree)], dtype=torch.int32)
_, pyramid, prefix = scan_octrees(octree, length)
query_points = torch.tensor(
[[3, 2, 0], [3, 1, 1], [0, 0, 0], [3, 3, 3], [2, 2, 2], [1, 1, 1]],
device='cuda',
dtype=torch.short)
point_hierarchy = generate_points(octree, pyramid, prefix)
results = unbatched_query(octree, prefix, query_points, 2)
expected_results = torch.tensor([7, 6, 5, 8, -1, -1],
dtype=torch.long,
device='cuda')
assert torch.equal(point_hierarchy[results[:-2]], query_points[:-2])
assert torch.equal(expected_results, results)
def test_ambiguous_raytrace(self):
# TODO(ttakikawa):
# Since 0.10.0, the behaviour of raytracing exactly between voxels
# has been changed from no hits at all to hitting all adjacent voxels.
# This has numerical ramifications because it may cause instability / error
# in the estimation of optical thickness in the volume rendering process
# among other issues. However, we have found that this doesn't lead to any
# obvious visual errors, whereas the no hit case causes speckle noise.
# We will eventually do a more thorough analysis of the numerical consideration of this
# behaviour, but for now we choose to prevent obvious visual errors.
octree = torch.tensor([255], dtype=torch.uint8, device='cuda')
length = torch.tensor([1], dtype=torch.int32)
max_level, pyramids, exsum = scan_octrees(octree, length)
point_hierarchy = generate_points(octree, pyramids, exsum)
origin = torch.tensor([[0., 0., 3.], [3., 3., 3.]],
dtype=torch.float,
device='cuda')
direction = torch.tensor(
[[0., 0., -1.], [-1. / 3., -1. / 3., -1. / 3.]],
dtype=torch.float,
device='cuda')
ridx, pidx, depth = unbatched_raytrace(octree,
point_hierarchy,
pyramids[0],
exsum,
origin,
direction,
1,
return_depth=True)
expected_nuggets = torch.tensor(
[[0, 2], [0, 1], [0, 4], [0, 6], [0, 3], [0, 5], [0, 8], [0, 7],
[1, 8], [1, 1]],
device='cuda',
dtype=torch.int)
assert torch.equal(ridx, expected_nuggets[..., 0])
assert torch.equal(pidx, expected_nuggets[..., 1])
def octree_to_spc(octree):
lengths = torch.tensor([len(octree)], dtype=torch.int32)
_, pyramid, prefix = spc_ops.scan_octrees(octree, lengths)
points = spc_ops.generate_points(octree, pyramid, prefix)
pyramid = pyramid[0]
return points, pyramid, prefix
def max_level_pyramids_exsum(self, octree, length):
return scan_octrees(octree, length)
def max_level_pyramids_exsum(self, octrees, lengths):
return spc.scan_octrees(octrees, lengths)