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)