"pytorch/run_text8_large.sh" did not exist on "c6b06f76bca1cda1cbce92d303c1855bd4224085"
points_in_boxes.py 3.72 KB
Newer Older
wuyuefeng's avatar
wuyuefeng committed
1
2
3
4
5
6
import torch

from . import roiaware_pool3d_ext


def points_in_boxes_gpu(points, boxes):
wuyuefeng's avatar
wuyuefeng committed
7
    """Find points that are in boxes (CUDA)
wuyuefeng's avatar
wuyuefeng committed
8

wuyuefeng's avatar
wuyuefeng committed
9
10
11
12
13
    Args:
        points (torch.Tensor): [B, M, 3], [x, y, z] in LiDAR coordinate
        boxes (torch.Tensor): [B, T, 7],
            num_valid_boxes <= T, [x, y, z, w, l, h, ry] in LiDAR coordinate,
            (x, y, z) is the bottom center
wuyuefeng's avatar
wuyuefeng committed
14

wuyuefeng's avatar
wuyuefeng committed
15
16
17
18
19
20
21
22
23
    Returns:
        box_idxs_of_pts (torch.Tensor): (B, M), default background = -1
    """
    assert boxes.shape[0] == points.shape[0]
    assert boxes.shape[2] == 7
    batch_size, num_points, _ = points.shape

    box_idxs_of_pts = points.new_zeros((batch_size, num_points),
                                       dtype=torch.int).fill_(-1)
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38

    # 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)

wuyuefeng's avatar
wuyuefeng committed
39
40
41
42
43
44
45
46
    roiaware_pool3d_ext.points_in_boxes_gpu(boxes.contiguous(),
                                            points.contiguous(),
                                            box_idxs_of_pts)

    return box_idxs_of_pts


def points_in_boxes_cpu(points, boxes):
wuyuefeng's avatar
wuyuefeng committed
47
48
    """Find points that are in boxes (CPU)

zhangwenwei's avatar
zhangwenwei committed
49
50
    Note:
        Currently, the output of this function is different from that of
wuyuefeng's avatar
wuyuefeng committed
51
        points_in_boxes_gpu.
wuyuefeng's avatar
wuyuefeng committed
52

wuyuefeng's avatar
wuyuefeng committed
53
54
55
56
    Args:
        points (torch.Tensor): [npoints, 3]
        boxes (torch.Tensor): [N, 7], in LiDAR coordinate,
            (x, y, z) is the bottom center
wuyuefeng's avatar
wuyuefeng committed
57

wuyuefeng's avatar
wuyuefeng committed
58
59
60
    Returns:
        point_indices (torch.Tensor): (N, npoints)
    """
wuyuefeng's avatar
wuyuefeng committed
61
    # TODO: Refactor this function as a CPU version of points_in_boxes_gpu
wuyuefeng's avatar
wuyuefeng committed
62
63
64
65
66
67
68
69
70
71
    assert boxes.shape[1] == 7
    assert points.shape[1] == 3

    point_indices = points.new_zeros((boxes.shape[0], points.shape[0]),
                                     dtype=torch.int)
    roiaware_pool3d_ext.points_in_boxes_cpu(boxes.float().contiguous(),
                                            points.float().contiguous(),
                                            point_indices)

    return point_indices
wuyuefeng's avatar
Votenet  
wuyuefeng committed
72
73
74
75
76
77
78
79
80


def points_in_boxes_batch(points, boxes):
    """Find points that are in boxes (CUDA)

    Args:
        points (torch.Tensor): [B, M, 3], [x, y, z] in LiDAR coordinate
        boxes (torch.Tensor): [B, T, 7],
            num_valid_boxes <= T, [x, y, z, w, l, h, ry] in LiDAR coordinate,
wuyuefeng's avatar
wuyuefeng committed
81
            (x, y, z) is the bottom center.
wuyuefeng's avatar
Votenet  
wuyuefeng committed
82
83
84
85
86
87
88
89
90
91
92

    Returns:
        box_idxs_of_pts (torch.Tensor): (B, M, T), default background = 0
    """
    assert boxes.shape[0] == points.shape[0]
    assert boxes.shape[2] == 7
    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)
93
94
95
96
97
98
99
100

    # 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)

wuyuefeng's avatar
Votenet  
wuyuefeng committed
101
102
103
104
105
    roiaware_pool3d_ext.points_in_boxes_batch(boxes.contiguous(),
                                              points.contiguous(),
                                              box_idxs_of_pts)

    return box_idxs_of_pts