(grid_coord, batch=None, depth=16, order="z")
| 7 | |
| 8 | @torch.inference_mode() |
| 9 | def encode(grid_coord, batch=None, depth=16, order="z"): |
| 10 | assert order in {"z", "z-trans", "hilbert", "hilbert-trans"} |
| 11 | if order == "z": |
| 12 | code = z_order_encode(grid_coord, depth=depth) |
| 13 | elif order == "z-trans": |
| 14 | code = z_order_encode(grid_coord[:, [1, 0, 2]], depth=depth) |
| 15 | elif order == "hilbert": |
| 16 | code = hilbert_encode(grid_coord, depth=depth) |
| 17 | elif order == "hilbert-trans": |
| 18 | code = hilbert_encode(grid_coord[:, [1, 0, 2]], depth=depth) |
| 19 | else: |
| 20 | raise NotImplementedError |
| 21 | if batch is not None: |
| 22 | batch = batch.long() |
| 23 | code = batch << depth * 3 | code |
| 24 | return code |
| 25 | |
| 26 | |
| 27 | @torch.inference_mode() |
no test coverage detected