File size: 3,439 Bytes
c6535db | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 | # Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved
import logging
import numpy as np
import torch
from sam3.perflib.masks_ops import mask_iou
try:
from torch_generic_nms import generic_nms as generic_nms_cuda
GENERIC_NMS_AVAILABLE = True
except ImportError:
logging.debug(
"Falling back to triton or CPU mask NMS implementation -- please install `torch_generic_nms` via\n\t"
'pip uninstall -y torch_generic_nms; TORCH_CUDA_ARCH_LIST="8.0 9.0" pip install git+https://github.com/ronghanghu/torch_generic_nms'
)
GENERIC_NMS_AVAILABLE = False
def nms_masks(
pred_probs: torch.Tensor,
pred_masks: torch.Tensor,
prob_threshold: float,
iou_threshold: float,
) -> torch.Tensor:
"""
Args:
- pred_probs: (num_det,) float Tensor, containing the score (probability) of each detection
- pred_masks: (num_det, H_mask, W_mask) float Tensor, containing the binary segmentation mask of each detection
- prob_threshold: float, score threshold to prefilter detections (NMS is performed on detections above threshold)
- iou_threshold: float, mask IoU threshold for NMS
Returns:
- keep: (num_det,) bool Tensor, indicating whether each detection is kept after score thresholding + NMS
"""
# prefilter the detections with prob_threshold ("valid" are those above prob_threshold)
is_valid = pred_probs > prob_threshold # (num_det,)
probs = pred_probs[is_valid] # (num_valid,)
masks_binary = pred_masks[is_valid] > 0 # (num_valid, H_mask, W_mask)
if probs.numel() == 0:
return is_valid # no valid detection, return empty keep mask
ious = mask_iou(masks_binary, masks_binary) # (num_valid, num_valid)
kept_inds = generic_nms(ious, probs, iou_threshold)
# valid_inds are the indices among `probs` of valid detections before NMS (or -1 for invalid)
valid_inds = torch.where(is_valid, is_valid.cumsum(dim=0) - 1, -1) # (num_det,)
keep = torch.isin(valid_inds, kept_inds) # (num_det,)
return keep
def generic_nms(
ious: torch.Tensor, scores: torch.Tensor, iou_threshold=0.5
) -> torch.Tensor:
"""A generic version of `torchvision.ops.nms` that takes a pairwise IoU matrix."""
assert ious.dim() == 2 and ious.size(0) == ious.size(1)
assert scores.dim() == 1 and scores.size(0) == ious.size(0)
if ious.is_cuda:
if GENERIC_NMS_AVAILABLE:
return generic_nms_cuda(ious, scores, iou_threshold, use_iou_matrix=True)
else:
from sam3.perflib.triton.nms import nms_triton
return nms_triton(ious, scores, iou_threshold)
return generic_nms_cpu(ious, scores, iou_threshold)
def generic_nms_cpu(
ious: torch.Tensor, scores: torch.Tensor, iou_threshold=0.5
) -> torch.Tensor:
"""
A generic version of `torchvision.ops.nms` that takes a pairwise IoU matrix. (CPU implementation
based on https://github.com/jwyang/faster-rcnn.pytorch/blob/master/lib/model/nms/nms_cpu.py)
"""
ious_np = ious.float().detach().cpu().numpy()
scores_np = scores.float().detach().cpu().numpy()
order = scores_np.argsort()[::-1]
kept_inds = []
while order.size > 0:
i = order.item(0)
kept_inds.append(i)
inds = np.where(ious_np[i, order[1:]] <= iou_threshold)[0]
order = order[inds + 1]
return torch.tensor(kept_inds, dtype=torch.int64, device=scores.device)
|