Operations for SPC aggregation and subdivision about kaolin

[charlesloop]

Hi,

Glad to hear you are finding the framework useful. I think I know what you want, but not quite sure what you want the API to look like. If you could provide more context and a code sample that would be helpful.

In the mean time, I think the following may be helpful:

def subdivide(points_in):
    #create level 1 (2x2x2) grid
    mort = torch.arange(8, dtype=torch.long, device='cuda')
    pts0 = kaolin.ops.spc.morton_to_points(mort).reshape(-1)
    #double input grid, add level 1 grid
    pts1 = (2*points_in).repeat(1, 8)
    return (pts1 + pts0).reshape(-1,3)

This takes as input a morton ordered tensor of points (voxels) of shape (3xN int16) and returns the corresponding tensor of 8 to 1 subdivided voxels of shape (3x(8*n)). For a given Spc, you would need to input the points from some level of the point hierarchy

subdiv_points = subdivide(point_hierarchies[pyramids[0,1,level]:pyramids[0,1,level+1]])

You could then create the octree for this by

subdiv_octree = kaolin.ops.spc.unbatched_points_octree(subdiv_points)

Admittedly, this is somewhat low level but I think close to what you want. LMK if you have thoughts for a higher level API, or if I'm way off the mark.

Good luck!

[akanimax]

@charlesloop,

Thanks a lot for this code snippet, this is exactly what I was trying to achieve for the subdivision part. Regarding a higher-level API, I was thinking something like this:
kaolin.ops.spc.subdivide(spc: SPC, mode:str = "random_zeros", subdiv_indices:Optional[Tensor]= None) -> SPC

Where, the SPC can have a batch > 1 and the subdivision routine will loop over all the points in the SPC batchwise to perform the subdivision. The mode can be one of three:

1. "random_zeros" puts the feature randomly in one of the 8 cells in the next hierarchy.
2. "zeros" puts the features at the subdiv_indices for the next level of hierarchy (rest locations being zeros). The subdiv_indices can be provided by the aggregator op similar to pooling and unpooling operations in 2D.
3. "nearest" copies the feature for all cells 8 times in the next hierarchy.

Nevertheless, thanks a lot for your help.

Cheers!
@akanimax

[akanimax]

Thanks again for your help with this.
Cheers!