in utils/box_ops3d.py [0:0]
def generalized_box3d_iou_tensor(corners1: torch.Tensor, corners2: torch.Tensor, nums_k2: torch.Tensor, rotated_boxes: bool = True,
return_inter_vols_only: bool = False, no_grad: bool = False):
"""
Input:
corners1: torch Tensor (B, K1, 8, 3), assume up direction is negative Y
corners2: torch Tensor (B, K2, 8, 3), assume up direction is negative Y
Assumes that the box is only rotated along Z direction
Returns:
B x K1 x K2 matrix of generalized IOU by approximating the boxes to be axis aligned
The return IOU is differentiable
"""
assert len(corners1.shape) == 4
assert len(corners2.shape)== 4
assert corners1.shape[2] == 8
assert corners1.shape[3] == 3
assert corners1.shape[0] == corners2.shape[0]
assert corners1.shape[2] == corners2.shape[2]
assert corners1.shape[3] == corners2.shape[3]
B, K1 = corners1.shape[0], corners1.shape[1]
_, K2 = corners2.shape[0], corners2.shape[1]
# # box height. Y is negative, so max is torch.min
ymax = torch.min(corners1[:, :, 0,1][:, :, None], corners2[:, :, 0,1][:, None, :])
ymin = torch.max(corners1[:, :, 4,1][:, :, None], corners2[:, :, 4,1][:, None, :])
height = (ymax - ymin).clamp(min=0)
EPS = 1e-8
idx = torch.arange(start=3, end=-1, step=-1, device=corners1.device)
idx2 = torch.tensor([0,2], dtype=torch.int64, device=corners1.device)
rect1 = corners1[:, :, idx, :]
rect2 = corners2[:, :, idx, :]
rect1 = rect1[:, :, :, idx2]
rect2 = rect2[:, :, :, idx2]
lt = torch.max(rect1[:, :, 1][:, :, None, :], rect2[:, :, 1][:, None, : ,:])
rb = torch.min(rect1[:, :, 3][:, :, None, :], rect2[:, :, 3][:, None, : ,:])
wh = (rb - lt).clamp(min=0)
non_rot_inter_areas = wh[:, :, :, 0] * wh[:, :, :, 1]
non_rot_inter_areas = non_rot_inter_areas.view(B, K1, K2)
if nums_k2 is not None:
for b in range(B):
non_rot_inter_areas[b, :, nums_k2[b]:] = 0
enclosing_vols = enclosing_box3d_vol(corners1, corners2)
# vols of boxes
vols1 = box3d_vol_tensor(corners1).clamp(min=EPS)
vols2 = box3d_vol_tensor(corners2).clamp(min=EPS)
sum_vols = vols1[:, :, None] + vols2[:, None, :]
# filter malformed boxes
good_boxes = (enclosing_vols > 2*EPS) * (sum_vols > 4*EPS)
if rotated_boxes:
inter_areas = torch.zeros((B, K1, K2), dtype=torch.float32)
rect1 = rect1.cpu()
rect2 = rect2.cpu()
nums_k2_np = to_list_1d(nums_k2)
non_rot_inter_areas_np = to_list_3d(non_rot_inter_areas)
for b in range(B):
for k1 in range(K1):
for k2 in range(K2):
if nums_k2 is not None and k2 >= nums_k2_np[b]:
break
if non_rot_inter_areas_np[b][k1][k2] == 0:
continue
##### compute volume of intersection
# inter = polygon_clip(rect1[b, k1], rect2[b, k2])
inter = polygon_clip_unnest(rect1[b, k1], rect2[b, k2])
# if inter is None:
# if len(inter) == 0:
# # area = torch.zeros(1, dtype=torch.float32, device=inter_areas.device).squeeze()
# # area = 0
# continue
# else:
if len(inter) > 0:
# inter = torch.stack(inter)
# xs = inter[:, 0]
# ys = inter[:, 1]
xs = torch.stack([x[0] for x in inter])
ys = torch.stack([x[1] for x in inter])
# area = poly_area_tensor(xs, ys)
inter_areas[b,k1,k2] = torch.abs(torch.dot(xs,torch.roll(ys,1))-torch.dot(ys,torch.roll(xs,1)))
inter_areas.mul_(0.5)
else:
inter_areas = non_rot_inter_areas
inter_areas = inter_areas.to(corners1.device)
### gIOU = iou - (1 - sum_vols/enclose_vol)
inter_vols = inter_areas * height
if return_inter_vols_only:
return inter_vols
union_vols = (sum_vols - inter_vols).clamp(min=EPS)
ious = inter_vols / union_vols
giou_second_term = - (1 - union_vols / enclosing_vols)
gious = ious + giou_second_term
gious *= good_boxes
if nums_k2 is not None:
mask = torch.zeros((B, K1, K2), device=height.device, dtype=torch.float32)
for b in range(B):
mask[b,:,:nums_k2[b]] = 1
gious *= mask
return gious