def boxes_iou3d_gpu(boxes_a, boxes_b):
"""
:param boxes_a: (N, 7) [x, y, z, h, w, l, ry]
:param boxes_b: (M, 7) [x, y, z, h, w, l, ry]
:return:
ans_iou: (M, N)
"""
boxes_a_bev = kitti_utils.boxes3d_to_bev_torch(boxes_a)
boxes_b_bev = kitti_utils.boxes3d_to_bev_torch(boxes_b)
# bev overlap
overlaps_bev = torch.cuda.FloatTensor(
torch.Size((boxes_a.shape[0], boxes_b.shape[0]))).zero_() # (N, M)
iou3d_cuda.boxes_overlap_bev_gpu(boxes_a_bev.contiguous(),
boxes_b_bev.contiguous(), overlaps_bev)
# height overlap
boxes_a_height_min = (boxes_a[:, 1] - boxes_a[:, 3]).view(-1, 1)
boxes_a_height_max = boxes_a[:, 1].view(-1, 1)
boxes_b_height_min = (boxes_b[:, 1] - boxes_b[:, 3]).view(1, -1)
boxes_b_height_max = boxes_b[:, 1].view(1, -1)
max_of_min = torch.max(boxes_a_height_min, boxes_b_height_min)
min_of_max = torch.min(boxes_a_height_max, boxes_b_height_max)
overlaps_h = torch.clamp(min_of_max - max_of_min, min=0)
# 3d iou
overlaps_3d = overlaps_bev * overlaps_h
vol_a = (boxes_a[:, 3] * boxes_a[:, 4] * boxes_a[:, 5]).view(-1, 1)
vol_b = (boxes_b[:, 3] * boxes_b[:, 4] * boxes_b[:, 5]).view(1, -1)
iou3d = overlaps_3d / torch.clamp(vol_a + vol_b - overlaps_3d, min=1e-7)
return iou3d
def boxes_iou3d_gpu(boxes_a,
boxes_b,
box_mode='wlh',
rect=False,
need_bev=False):
"""
Input (torch):
boxes_a: (N, 7) [x, y, z, h, w, l, ry], torch tensor with type float32
boxes_b: (M, 7) [x, y, z, h, w, l, ry], torch tensor with type float32
rect: True/False means boxes in camera/velodyne coord system.
Output:
iou_3d: (N, M)
"""
w_index, l_index, h_index = box_mode.index('w') + 3, box_mode.index(
'l') + 3, box_mode.index('h') + 3
boxes_a_bev = utils.boxes3d_to_bev_torch(boxes_a, box_mode, rect)
boxes_b_bev = utils.boxes3d_to_bev_torch(boxes_b, box_mode, rect)
# bev overlap
overlaps_bev = torch.cuda.FloatTensor(
torch.Size((boxes_a.shape[0], boxes_b.shape[0]))).zero_() # (N, M)
iou3d_cuda.boxes_overlap_bev_gpu(boxes_a_bev.contiguous(),
boxes_b_bev.contiguous(), overlaps_bev)
# bev iou
area_a = (boxes_a[:, w_index] * boxes_a[:, l_index]).view(-1, 1) # (N, 1)
area_b = (boxes_b[:, w_index] * boxes_b[:, l_index]).view(
1, -1) # (1, M) -> broadcast (N, M)
iou_bev = overlaps_bev / torch.clamp(area_a + area_b - overlaps_bev,
min=1e-7)
# height overlap
if rect:
boxes_a_height_min = (boxes_a[:, 1] - boxes_a[:, h_index]).view(
-1, 1) # y - h
boxes_a_height_max = boxes_a[:, 1].view(-1, 1) # y
boxes_b_height_min = (boxes_b[:, 1] - boxes_b[:, h_index]).view(1, -1)
boxes_b_height_max = boxes_b[:, 1].view(1, -1)
else:
# todo: notice if (x, y, z) is the real center
half_h_a = boxes_a[:, h_index] / 2.0
half_h_b = boxes_b[:, h_index] / 2.0
boxes_a_height_min = (boxes_a[:, 2] - half_h_a).view(
-1, 1) # z - h/2, (N, 1)
boxes_a_height_max = (boxes_a[:, 2] + half_h_a).view(
-1, 1) # z + h/2, (N, 1)
boxes_b_height_min = (boxes_b[:, 2] - half_h_b).view(1, -1)
boxes_b_height_max = (boxes_b[:, 2] + half_h_b).view(1, -1)
max_of_min = torch.max(boxes_a_height_min, boxes_b_height_min) # (N, M)
min_of_max = torch.min(boxes_a_height_max, boxes_b_height_max) # (N, M)
overlaps_h = torch.clamp(min_of_max - max_of_min, min=0) # (N, M)
# 3d iou
overlaps_3d = overlaps_bev * overlaps_h # broadcast: (N, M)
vol_a = (boxes_a[:, 3] * boxes_a[:, 4] * boxes_a[:, 5]).view(-1,
1) # (N, 1)
vol_b = (boxes_b[:, 3] * boxes_b[:, 4] * boxes_b[:, 5]).view(
1, -1) # (1, M) -> broadcast (N, M)
iou3d = overlaps_3d / torch.clamp(vol_a + vol_b - overlaps_3d, min=1e-7)
if need_bev:
return iou3d, iou_bev
return iou3d