Perform mask non-maximum suppression (NMS) on a set of masks based on their scores. Args: masks (torch.Tensor): has shape (num_masks, H, W) scores (torch.Tensor): The scores of the masks, has shape (num_masks,) iou_thr (float, optional): The threshold for IoU.
(masks, scores, iou_thr=0.7, score_thr=0.1, inner_thr=0.2, **kwargs)
| 216 | return result_keep |
| 217 | |
| 218 | def mask_nms(masks, scores, iou_thr=0.7, score_thr=0.1, inner_thr=0.2, **kwargs): |
| 219 | """ |
| 220 | Perform mask non-maximum suppression (NMS) on a set of masks based on their scores. |
| 221 | |
| 222 | Args: |
| 223 | masks (torch.Tensor): has shape (num_masks, H, W) |
| 224 | scores (torch.Tensor): The scores of the masks, has shape (num_masks,) |
| 225 | iou_thr (float, optional): The threshold for IoU. |
| 226 | score_thr (float, optional): The threshold for the mask scores. |
| 227 | inner_thr (float, optional): The threshold for the overlap rate. |
| 228 | **kwargs: Additional keyword arguments. |
| 229 | Returns: |
| 230 | selected_idx (torch.Tensor): A tensor representing the selected indices of the masks after NMS. |
| 231 | """ |
| 232 | |
| 233 | scores, idx = scores.sort(0, descending=True) |
| 234 | num_masks = idx.shape[0] |
| 235 | |
| 236 | masks_ord = masks[idx.view(-1), :] |
| 237 | masks_area = torch.sum(masks_ord, dim=(1, 2), dtype=torch.float) |
| 238 | |
| 239 | iou_matrix = torch.zeros((num_masks,) * 2, dtype=torch.float, device=masks.device) |
| 240 | inner_iou_matrix = torch.zeros((num_masks,) * 2, dtype=torch.float, device=masks.device) |
| 241 | for i in range(num_masks): |
| 242 | for j in range(i, num_masks): |
| 243 | intersection = torch.sum(torch.logical_and(masks_ord[i], masks_ord[j]), dtype=torch.float) |
| 244 | union = torch.sum(torch.logical_or(masks_ord[i], masks_ord[j]), dtype=torch.float) |
| 245 | iou = intersection / union |
| 246 | iou_matrix[i, j] = iou |
| 247 | # select mask pairs that may have a severe internal relationship |
| 248 | if intersection / masks_area[i] < 0.5 and intersection / masks_area[j] >= 0.85: |
| 249 | inner_iou = 1 - (intersection / masks_area[j]) * (intersection / masks_area[i]) |
| 250 | inner_iou_matrix[i, j] = inner_iou |
| 251 | if intersection / masks_area[i] >= 0.85 and intersection / masks_area[j] < 0.5: |
| 252 | inner_iou = 1 - (intersection / masks_area[j]) * (intersection / masks_area[i]) |
| 253 | inner_iou_matrix[j, i] = inner_iou |
| 254 | |
| 255 | iou_matrix.triu_(diagonal=1) |
| 256 | iou_max, _ = iou_matrix.max(dim=0) |
| 257 | inner_iou_matrix_u = torch.triu(inner_iou_matrix, diagonal=1) |
| 258 | inner_iou_max_u, _ = inner_iou_matrix_u.max(dim=0) |
| 259 | inner_iou_matrix_l = torch.tril(inner_iou_matrix, diagonal=1) |
| 260 | inner_iou_max_l, _ = inner_iou_matrix_l.max(dim=0) |
| 261 | |
| 262 | keep = iou_max <= iou_thr |
| 263 | keep_conf = scores > score_thr |
| 264 | keep_inner_u = inner_iou_max_u <= 1 - inner_thr |
| 265 | keep_inner_l = inner_iou_max_l <= 1 - inner_thr |
| 266 | |
| 267 | # If there are no masks with scores above threshold, the top 3 masks are selected |
| 268 | if keep_conf.sum() == 0: |
| 269 | index = scores.topk(3).indices |
| 270 | keep_conf[index, 0] = True |
| 271 | if keep_inner_u.sum() == 0: |
| 272 | index = scores.topk(3).indices |
| 273 | keep_inner_u[index, 0] = True |
| 274 | if keep_inner_l.sum() == 0: |
| 275 | index = scores.topk(3).indices |