import torch import spconv.pytorch as spconv try: import ocnn except ImportError: ocnn = None from addict import Dict from pointcept.models.utils.serialization import encode, decode from pointcept.models.utils import offset2batch, batch2offset class Point(Dict): """ Point Structure of Pointcept A Point (point cloud) in Pointcept is a dictionary that contains various properties of a batched point cloud. The property with the following names have a specific definition as follows: - "coord": original coordinate of point cloud; - "grid_coord": grid coordinate for specific grid size (related to GridSampling); Point also support the following optional attributes: - "offset": if not exist, initialized as batch size is 1; - "batch": if not exist, initialized as batch size is 1; - "feat": feature of point cloud, default input of model; - "grid_size": Grid size of point cloud (related to GridSampling); (related to Serialization) - "serialized_depth": depth of serialization, 2 ** depth * grid_size describe the maximum of point cloud range; - "serialized_code": a list of serialization codes; - "serialized_order": a list of serialization order determined by code; - "serialized_inverse": a list of inverse mapping determined by code; (related to Sparsify: SpConv) - "sparse_shape": Sparse shape for Sparse Conv Tensor; - "sparse_conv_feat": SparseConvTensor init with information provide by Point; """ def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) # If one of "offset" or "batch" do not exist, generate by the existing one if "batch" not in self.keys() and "offset" in self.keys(): self["batch"] = offset2batch(self.offset) elif "offset" not in self.keys() and "batch" in self.keys(): self["offset"] = batch2offset(self.batch) def serialization(self, order="z", depth=None, shuffle_orders=False): """ Point Cloud Serialization relay on ["grid_coord" or "coord" + "grid_size", "batch", "feat"] """ assert "batch" in self.keys() if "grid_coord" not in self.keys(): # if you don't want to operate GridSampling in data augmentation, # please add the following augmentation into your pipline: # dict(type="Copy", keys_dict={"grid_size": 0.01}), # (adjust `grid_size` to what your want) assert {"grid_size", "coord"}.issubset(self.keys()) self["grid_coord"] = torch.div( self.coord - self.coord.min(0)[0], self.grid_size, rounding_mode="trunc" ).int() if depth is None: # Adaptive measure the depth of serialization cube (length = 2 ^ depth) depth = int(self.grid_coord.max()).bit_length() self["serialized_depth"] = depth # Maximum bit length for serialization code is 63 (int64) assert depth * 3 + len(self.offset).bit_length() <= 63 # Here we follow OCNN and set the depth limitation to 16 (48bit) for the point position. # Although depth is limited to less than 16, we can encode a 655.36^3 (2^16 * 0.01) meter^3 # cube with a grid size of 0.01 meter. We consider it is enough for the current stage. # We can unlock the limitation by optimizing the z-order encoding function if necessary. assert depth <= 16 # The serialization codes are arranged as following structures: # [Order1 ([n]), # Order2 ([n]), # ... # OrderN ([n])] (k, n) code = [ encode(self.grid_coord, self.batch, depth, order=order_) for order_ in order ] code = torch.stack(code) order = torch.argsort(code) inverse = torch.zeros_like(order).scatter_( dim=1, index=order, src=torch.arange(0, code.shape[1], device=order.device).repeat( code.shape[0], 1 ), ) if shuffle_orders: perm = torch.randperm(code.shape[0]) code = code[perm] order = order[perm] inverse = inverse[perm] self["serialized_code"] = code self["serialized_order"] = order self["serialized_inverse"] = inverse def sparsify(self, pad=96): """ Point Cloud Serialization Point cloud is sparse, here we use "sparsify" to specifically refer to preparing "spconv.SparseConvTensor" for SpConv. relay on ["grid_coord" or "coord" + "grid_size", "batch", "feat"] pad: padding sparse for sparse shape. """ assert {"feat", "batch"}.issubset(self.keys()) if "grid_coord" not in self.keys(): # if you don't want to operate GridSampling in data augmentation, # please add the following augmentation into your pipline: # dict(type="Copy", keys_dict={"grid_size": 0.01}), # (adjust `grid_size` to what your want) assert {"grid_size", "coord"}.issubset(self.keys()) self["grid_coord"] = torch.div( self.coord - self.coord.min(0)[0], self.grid_size, rounding_mode="trunc" ).int() if "sparse_shape" in self.keys(): sparse_shape = self.sparse_shape else: sparse_shape = torch.add( torch.max(self.grid_coord, dim=0).values, pad ).tolist() sparse_conv_feat = spconv.SparseConvTensor( features=self.feat, indices=torch.cat( [self.batch.unsqueeze(-1).int(), self.grid_coord.int()], dim=1 ).contiguous(), spatial_shape=sparse_shape, batch_size=self.batch[-1].tolist() + 1, ) self["sparse_shape"] = sparse_shape self["sparse_conv_feat"] = sparse_conv_feat def octreetization(self, depth=None, full_depth=None): """ Point Cloud Octreelization Generate octree with OCNN relay on ["grid_coord", "batch", "feat"] """ assert ( ocnn is not None ), "Please follow https://github.com/octree-nn/ocnn-pytorch install ocnn." assert {"grid_coord", "feat", "batch"}.issubset(self.keys()) # add 1 to make grid space support shift order if depth is None: if "depth" in self.keys(): depth = self.depth else: depth = int(self.grid_coord.max() + 1).bit_length() if full_depth is None: full_depth = 2 self["depth"] = depth assert depth <= 16 # maximum in ocnn # [0, 2**depth] -> [0, 2] -> [-1, 1] coord = self.grid_coord / 2 ** (self.depth - 1) - 1.0 point = ocnn.octree.Points( points=coord, features=self.feat, batch_id=self.batch.unsqueeze(-1), batch_size=self.batch[-1] + 1, ) octree = ocnn.octree.Octree( depth=depth, full_depth=full_depth, batch_size=self.batch[-1] + 1, device=coord.device, ) octree.build_octree(point) octree.construct_all_neigh() self["octree"] = octree