3d nms for aligned boxes. Args: boxes (torch.Tensor): Aligned box with shape [n, 6]. scores (torch.Tensor): Scores of each box. classes (torch.Tensor): Class of each box. thresh (float): Iou threshold for nms. Returns: torch.Tensor: Indices of select
(boxes, scores, classes, thresh)
| 89 | |
| 90 | |
| 91 | def aligned_3d_nms(boxes, scores, classes, thresh): |
| 92 | """3d nms for aligned boxes. |
| 93 | |
| 94 | Args: |
| 95 | boxes (torch.Tensor): Aligned box with shape [n, 6]. |
| 96 | scores (torch.Tensor): Scores of each box. |
| 97 | classes (torch.Tensor): Class of each box. |
| 98 | thresh (float): Iou threshold for nms. |
| 99 | |
| 100 | Returns: |
| 101 | torch.Tensor: Indices of selected boxes. |
| 102 | """ |
| 103 | x1 = boxes[:, 0] |
| 104 | y1 = boxes[:, 1] |
| 105 | z1 = boxes[:, 2] |
| 106 | x2 = boxes[:, 3] |
| 107 | y2 = boxes[:, 4] |
| 108 | z2 = boxes[:, 5] |
| 109 | area = (x2 - x1) * (y2 - y1) * (z2 - z1) |
| 110 | zero = boxes.new_zeros(1, ) |
| 111 | |
| 112 | score_sorted = torch.argsort(scores) |
| 113 | pick = [] |
| 114 | while (score_sorted.shape[0] != 0): |
| 115 | last = score_sorted.shape[0] |
| 116 | i = score_sorted[-1] |
| 117 | pick.append(i) |
| 118 | |
| 119 | xx1 = torch.max(x1[i], x1[score_sorted[:last - 1]]) |
| 120 | yy1 = torch.max(y1[i], y1[score_sorted[:last - 1]]) |
| 121 | zz1 = torch.max(z1[i], z1[score_sorted[:last - 1]]) |
| 122 | xx2 = torch.min(x2[i], x2[score_sorted[:last - 1]]) |
| 123 | yy2 = torch.min(y2[i], y2[score_sorted[:last - 1]]) |
| 124 | zz2 = torch.min(z2[i], z2[score_sorted[:last - 1]]) |
| 125 | classes1 = classes[i] |
| 126 | classes2 = classes[score_sorted[:last - 1]] |
| 127 | inter_l = torch.max(zero, xx2 - xx1) |
| 128 | inter_w = torch.max(zero, yy2 - yy1) |
| 129 | inter_h = torch.max(zero, zz2 - zz1) |
| 130 | |
| 131 | inter = inter_l * inter_w * inter_h |
| 132 | iou = inter / (area[i] + area[score_sorted[:last - 1]] - inter) |
| 133 | iou = iou * (classes1 == classes2).float() |
| 134 | score_sorted = score_sorted[torch.nonzero( |
| 135 | iou <= thresh, as_tuple=False).flatten()] |
| 136 | |
| 137 | indices = boxes.new_tensor(pick, dtype=torch.long) |
| 138 | return indices |
| 139 | |
| 140 | |
| 141 | @numba.jit(nopython=True) |
no outgoing calls