Point Cloud Serialization relay on ["grid_coord" or "coord" + "grid_size", "batch", "feat"]
(self, order="z", depth=None, shuffle_orders=False)
| 79 | self["offset"] = batch2offset(self.batch) |
| 80 | |
| 81 | def serialization(self, order="z", depth=None, shuffle_orders=False): |
| 82 | """ |
| 83 | Point Cloud Serialization |
| 84 | |
| 85 | relay on ["grid_coord" or "coord" + "grid_size", "batch", "feat"] |
| 86 | """ |
| 87 | assert "batch" in self.keys() |
| 88 | if "grid_coord" not in self.keys(): |
| 89 | # if you don't want to operate GridSampling in data augmentation, |
| 90 | # please add the following augmentation into your pipline: |
| 91 | # dict(type="Copy", keys_dict={"grid_size": 0.01}), |
| 92 | # (adjust `grid_size` to what your want) |
| 93 | assert {"grid_size", "coord"}.issubset(self.keys()) |
| 94 | self["grid_coord"] = torch.div( |
| 95 | self.coord - self.coord.min(0)[0], self.grid_size, rounding_mode="trunc" |
| 96 | ).int() |
| 97 | |
| 98 | if depth is None: |
| 99 | # Adaptive measure the depth of serialization cube (length = 2 ^ depth) |
| 100 | depth = int(self.grid_coord.max()).bit_length() |
| 101 | self["serialized_depth"] = depth |
| 102 | # Maximum bit length for serialization code is 63 (int64) |
| 103 | assert depth * 3 + len(self.offset).bit_length() <= 63 |
| 104 | # Here we follow OCNN and set the depth limitation to 16 (48bit) for the point position. |
| 105 | # Although depth is limited to less than 16, we can encode a 655.36^3 (2^16 * 0.01) meter^3 |
| 106 | # cube with a grid size of 0.01 meter. We consider it is enough for the current stage. |
| 107 | # We can unlock the limitation by optimizing the z-order encoding function if necessary. |
| 108 | assert depth <= 16 |
| 109 | |
| 110 | # The serialization codes are arranged as following structures: |
| 111 | # [Order1 ([n]), |
| 112 | # Order2 ([n]), |
| 113 | # ... |
| 114 | # OrderN ([n])] (k, n) |
| 115 | code = [ |
| 116 | encode(self.grid_coord, self.batch, depth, order=order_) for order_ in order |
| 117 | ] |
| 118 | code = torch.stack(code) |
| 119 | order = torch.argsort(code) |
| 120 | inverse = torch.zeros_like(order).scatter_( |
| 121 | dim=1, |
| 122 | index=order, |
| 123 | src=torch.arange(0, code.shape[1], device=order.device).repeat( |
| 124 | code.shape[0], 1 |
| 125 | ), |
| 126 | ) |
| 127 | |
| 128 | if shuffle_orders: |
| 129 | perm = torch.randperm(code.shape[0]) |
| 130 | code = code[perm] |
| 131 | order = order[perm] |
| 132 | inverse = inverse[perm] |
| 133 | |
| 134 | self["serialized_code"] = code |
| 135 | self["serialized_order"] = order |
| 136 | self["serialized_inverse"] = inverse |
| 137 | |
| 138 | def sparsify(self, pad=96): |