Skip to content

GitLab

Switch branch/tag
Find file Normal viewHistoryPermalink
test_ball_query.py 5.19 KB
Newer Older
Zaida Zhou's avatar
Add copyright pre-commit-hook (#1742)  
Zaida Zhou committed
1 
# Copyright (c) OpenMMLab. All rights reserved.
dingchang's avatar
[Feature] Add ballquery op from mmdet3d (#1332)  
dingchang committed
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55 
import pytest
import torch

from mmcv.ops import ball_query


@pytest.mark.skipif(
    not torch.cuda.is_available(), reason='requires CUDA support')
def test_ball_query():
    new_xyz = torch.tensor([[[-0.0740, 1.3147, -1.3625],
                             [-2.2769, 2.7817, -0.2334],
                             [-0.4003, 2.4666, -0.5116],
                             [-0.0740, 1.3147, -1.3625],
                             [-0.0740, 1.3147, -1.3625]],
                            [[-2.0289, 2.4952, -0.1708],
                             [-2.0668, 6.0278, -0.4875],
                             [0.4066, 1.4211, -0.2947],
                             [-2.0289, 2.4952, -0.1708],
                             [-2.0289, 2.4952, -0.1708]]]).cuda()

    xyz = torch.tensor([[[-0.0740, 1.3147, -1.3625], [0.5555, 1.0399, -1.3634],
                         [-0.4003, 2.4666,
                          -0.5116], [-0.5251, 2.4379, -0.8466],
                         [-0.9691, 1.1418,
                          -1.3733], [-0.2232, 0.9561, -1.3626],
                         [-2.2769, 2.7817, -0.2334],
                         [-0.2822, 1.3192, -1.3645], [0.1533, 1.5024, -1.0432],
                         [0.4917, 1.1529, -1.3496]],
                        [[-2.0289, 2.4952,
                          -0.1708], [-0.7188, 0.9956, -0.5096],
                         [-2.0668, 6.0278, -0.4875], [-1.9304, 3.3092, 0.6610],
                         [0.0949, 1.4332, 0.3140], [-1.2879, 2.0008, -0.7791],
                         [-0.7252, 0.9611, -0.6371], [0.4066, 1.4211, -0.2947],
                         [0.3220, 1.4447, 0.3548], [-0.9744, 2.3856,
                                                    -1.2000]]]).cuda()

    idx = ball_query(0, 0.2, 5, xyz, new_xyz)
    expected_idx = torch.tensor([[[0, 0, 0, 0, 0], [6, 6, 6, 6, 6],
                                  [2, 2, 2, 2, 2], [0, 0, 0, 0, 0],
                                  [0, 0, 0, 0, 0]],
                                 [[0, 0, 0, 0, 0], [2, 2, 2, 2, 2],
                                  [7, 7, 7, 7, 7], [0, 0, 0, 0, 0],
                                  [0, 0, 0, 0, 0]]]).cuda()
    assert torch.all(idx == expected_idx)

    # test dilated ball query
    idx = ball_query(0.2, 0.4, 5, xyz, new_xyz)
    expected_idx = torch.tensor([[[0, 5, 7, 0, 0], [6, 6, 6, 6, 6],
                                  [2, 3, 2, 2, 2], [0, 5, 7, 0, 0],
                                  [0, 5, 7, 0, 0]],
                                 [[0, 0, 0, 0, 0], [2, 2, 2, 2, 2],
                                  [7, 7, 7, 7, 7], [0, 0, 0, 0, 0],
                                  [0, 0, 0, 0, 0]]]).cuda()
    assert torch.all(idx == expected_idx)
VVsssssk's avatar
[Features] Add stack ball query and stack group points ops (#2292)  
VVsssssk committed
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102 


@pytest.mark.skipif(
    not torch.cuda.is_available(), reason='requires CUDA support')
def test_stack_ball_query():
    new_xyz = torch.tensor([[-0.0740, 1.3147, -1.3625],
                            [-2.2769, 2.7817, -0.2334],
                            [-0.4003, 2.4666, -0.5116],
                            [-0.0740, 1.3147, -1.3625],
                            [-0.0740, 1.3147, -1.3625],
                            [-2.0289, 2.4952, -0.1708],
                            [-2.0668, 6.0278, -0.4875],
                            [0.4066, 1.4211, -0.2947],
                            [-2.0289, 2.4952, -0.1708],
                            [-2.0289, 2.4952, -0.1708]]).cuda()
    new_xyz_batch_cnt = torch.tensor([5, 5], dtype=torch.int32).cuda()
    xyz = torch.tensor([[-0.0740, 1.3147, -1.3625], [0.5555, 1.0399, -1.3634],
                        [-0.4003, 2.4666, -0.5116], [-0.5251, 2.4379, -0.8466],
                        [-0.9691, 1.1418, -1.3733], [-0.2232, 0.9561, -1.3626],
                        [-2.2769, 2.7817, -0.2334], [-0.2822, 1.3192, -1.3645],
                        [0.1533, 1.5024, -1.0432], [0.4917, 1.1529, -1.3496],
                        [-2.0289, 2.4952, -0.1708], [-0.7188, 0.9956, -0.5096],
                        [-2.0668, 6.0278, -0.4875], [-1.9304, 3.3092, 0.6610],
                        [0.0949, 1.4332, 0.3140], [-1.2879, 2.0008, -0.7791],
                        [-0.7252, 0.9611, -0.6371], [0.4066, 1.4211, -0.2947],
                        [0.3220, 1.4447, 0.3548], [-0.9744, 2.3856,
                                                   -1.2000]]).cuda()
    xyz_batch_cnt = torch.tensor([10, 10], dtype=torch.int32).cuda()
    idx = ball_query(0, 0.2, 5, xyz, new_xyz, xyz_batch_cnt, new_xyz_batch_cnt)
    expected_idx = torch.tensor([[0, 0, 0, 0, 0], [6, 6, 6, 6, 6],
                                 [2, 2, 2, 2, 2], [0, 0, 0, 0, 0],
                                 [0, 0, 0, 0, 0], [0, 0, 0, 0, 0],
                                 [2, 2, 2, 2, 2], [7, 7, 7, 7, 7],
                                 [0, 0, 0, 0, 0], [0, 0, 0, 0, 0]]).cuda()
    assert torch.all(idx == expected_idx)

    xyz = xyz.double()
    new_xyz = new_xyz.double()
    expected_idx = expected_idx.double()
    idx = ball_query(0, 0.2, 5, xyz, new_xyz, xyz_batch_cnt, new_xyz_batch_cnt)
    assert torch.all(idx == expected_idx)

    xyz = xyz.half()
    new_xyz = new_xyz.half()
    expected_idx = expected_idx.half()
    idx = ball_query(0, 0.2, 5, xyz, new_xyz, xyz_batch_cnt, new_xyz_batch_cnt)
    assert torch.all(idx == expected_idx)