| import torch |
| import vox2seq |
|
|
|
|
| if __name__ == "__main__": |
| RES = 256 |
| coords = torch.meshgrid(torch.arange(RES), torch.arange(RES), torch.arange(RES)) |
| coords = torch.stack(coords, dim=-1).reshape(-1, 3).int().cuda() |
| code_z_cuda = vox2seq.encode(coords, mode='z_order') |
| code_z_pytorch = vox2seq.pytorch.encode(coords, mode='z_order') |
| code_h_cuda = vox2seq.encode(coords, mode='hilbert') |
| code_h_pytorch = vox2seq.pytorch.encode(coords, mode='hilbert') |
| assert torch.equal(code_z_cuda, code_z_pytorch) |
| assert torch.equal(code_h_cuda, code_h_pytorch) |
|
|
| code = torch.arange(RES**3).int().cuda() |
| coords_z_cuda = vox2seq.decode(code, mode='z_order') |
| coords_z_pytorch = vox2seq.pytorch.decode(code, mode='z_order') |
| coords_h_cuda = vox2seq.decode(code, mode='hilbert') |
| coords_h_pytorch = vox2seq.pytorch.decode(code, mode='hilbert') |
| assert torch.equal(coords_z_cuda, coords_z_pytorch) |
| assert torch.equal(coords_h_cuda, coords_h_pytorch) |
|
|
| print("All tests passed.") |
|
|
|
|
