points_in_boxes.py 5.01 KB
Newer Older
1
# Copyright (c) OpenMMLab. All rights reserved.
wuyuefeng's avatar
wuyuefeng committed
2
3
4
5
6
import torch

from . import roiaware_pool3d_ext


7
8
def points_in_boxes_part(points, boxes):
    """Find the box in which each point is (CUDA).
wuyuefeng's avatar
wuyuefeng committed
9

wuyuefeng's avatar
wuyuefeng committed
10
    Args:
11
        points (torch.Tensor): [B, M, 3], [x, y, z] in LiDAR/DEPTH coordinate
wuyuefeng's avatar
wuyuefeng committed
12
        boxes (torch.Tensor): [B, T, 7],
13
            num_valid_boxes <= T, [x, y, z, x_size, y_size, z_size, rz] in
14
            LiDAR/DEPTH coordinate, (x, y, z) is the bottom center
wuyuefeng's avatar
wuyuefeng committed
15

wuyuefeng's avatar
wuyuefeng committed
16
17
18
    Returns:
        box_idxs_of_pts (torch.Tensor): (B, M), default background = -1
    """
19
    assert points.shape[0] == boxes.shape[0], \
20
        f'Points and boxes should have the same batch size, ' \
21
        f'got {points.shape[0]} and {boxes.shape[0]}'
22
23
24
25
26
27
    assert boxes.shape[2] == 7, \
        f'boxes dimension should be 7, ' \
        f'got unexpected shape {boxes.shape[2]}'
    assert points.shape[2] == 3, \
        f'points dimension should be 3, ' \
        f'got unexpected shape {points.shape[2]}'
wuyuefeng's avatar
wuyuefeng committed
28
29
30
31
    batch_size, num_points, _ = points.shape

    box_idxs_of_pts = points.new_zeros((batch_size, num_points),
                                       dtype=torch.int).fill_(-1)
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46

    # If manually put the tensor 'points' or 'boxes' on a device
    # which is not the current device, some temporary variables
    # will be created on the current device in the cuda op,
    # and the output will be incorrect.
    # Therefore, we force the current device to be the same
    # as the device of the tensors if it was not.
    # Please refer to https://github.com/open-mmlab/mmdetection3d/issues/305
    # for the incorrect output before the fix.
    points_device = points.get_device()
    assert points_device == boxes.get_device(), \
        'Points and boxes should be put on the same device'
    if torch.cuda.current_device() != points_device:
        torch.cuda.set_device(points_device)

47
48
49
    roiaware_pool3d_ext.points_in_boxes_part(boxes.contiguous(),
                                             points.contiguous(),
                                             box_idxs_of_pts)
wuyuefeng's avatar
wuyuefeng committed
50
51
52
53
54

    return box_idxs_of_pts


def points_in_boxes_cpu(points, boxes):
55
56
    """Find all boxes in which each point is (CPU). The CPU version of
    :meth:`points_in_boxes_all`.
wuyuefeng's avatar
wuyuefeng committed
57

wuyuefeng's avatar
wuyuefeng committed
58
    Args:
59
60
61
        points (torch.Tensor): [B, M, 3], [x, y, z] in
            LiDAR/DEPTH coordinate
        boxes (torch.Tensor): [B, T, 7],
62
            num_valid_boxes <= T, [x, y, z, x_size, y_size, z_size, rz],
63
            (x, y, z) is the bottom center.
wuyuefeng's avatar
wuyuefeng committed
64

wuyuefeng's avatar
wuyuefeng committed
65
    Returns:
66
        box_idxs_of_pts (torch.Tensor): (B, M, T), default background = 0.
wuyuefeng's avatar
wuyuefeng committed
67
    """
68
69
70
71
    assert points.shape[0] == boxes.shape[0], \
        f'Points and boxes should have the same batch size, ' \
        f'got {points.shape[0]} and {boxes.shape[0]}'
    assert boxes.shape[2] == 7, \
72
73
        f'boxes dimension should be 7, ' \
        f'got unexpected shape {boxes.shape[2]}'
74
    assert points.shape[2] == 3, \
75
76
        f'points dimension should be 3, ' \
        f'got unexpected shape {points.shape[2]}'
77
78
    batch_size, num_points, _ = points.shape
    num_boxes = boxes.shape[1]
wuyuefeng's avatar
wuyuefeng committed
79

80
    point_indices = points.new_zeros((batch_size, num_boxes, num_points),
wuyuefeng's avatar
wuyuefeng committed
81
                                     dtype=torch.int)
82
83
84
85
86
    for b in range(batch_size):
        roiaware_pool3d_ext.points_in_boxes_cpu(boxes[b].float().contiguous(),
                                                points[b].float().contiguous(),
                                                point_indices[b])
    point_indices = point_indices.transpose(1, 2)
wuyuefeng's avatar
wuyuefeng committed
87
88

    return point_indices
wuyuefeng's avatar
Votenet  
wuyuefeng committed
89
90


91
92
def points_in_boxes_all(points, boxes):
    """Find all boxes in which each point is (CUDA).
wuyuefeng's avatar
Votenet  
wuyuefeng committed
93
94

    Args:
95
        points (torch.Tensor): [B, M, 3], [x, y, z] in LiDAR/DEPTH coordinate
wuyuefeng's avatar
Votenet  
wuyuefeng committed
96
        boxes (torch.Tensor): [B, T, 7],
97
            num_valid_boxes <= T, [x, y, z, x_size, y_size, z_size, rz],
wuyuefeng's avatar
wuyuefeng committed
98
            (x, y, z) is the bottom center.
wuyuefeng's avatar
Votenet  
wuyuefeng committed
99
100

    Returns:
101
        box_idxs_of_pts (torch.Tensor): (B, M, T), default background = 0.
wuyuefeng's avatar
Votenet  
wuyuefeng committed
102
    """
103
104
105
106
107
108
109
110
111
    assert boxes.shape[0] == points.shape[0], \
        f'Points and boxes should have the same batch size, ' \
        f'got {boxes.shape[0]} and {boxes.shape[0]}'
    assert boxes.shape[2] == 7, \
        f'boxes dimension should be 7, ' \
        f'got unexpected shape {boxes.shape[2]}'
    assert points.shape[2] == 3, \
        f'points dimension should be 3, ' \
        f'got unexpected shape {points.shape[2]}'
wuyuefeng's avatar
Votenet  
wuyuefeng committed
112
113
114
115
116
    batch_size, num_points, _ = points.shape
    num_boxes = boxes.shape[1]

    box_idxs_of_pts = points.new_zeros((batch_size, num_points, num_boxes),
                                       dtype=torch.int).fill_(0)
117
118
119
120
121
122
123
124

    # Same reason as line 25-32
    points_device = points.get_device()
    assert points_device == boxes.get_device(), \
        'Points and boxes should be put on the same device'
    if torch.cuda.current_device() != points_device:
        torch.cuda.set_device(points_device)

125
126
127
    roiaware_pool3d_ext.points_in_boxes_all(boxes.contiguous(),
                                            points.contiguous(),
                                            box_idxs_of_pts)
wuyuefeng's avatar
Votenet  
wuyuefeng committed
128
129

    return box_idxs_of_pts