Non-maximum suppression

pip install torch torchvision

import torch
from torchvision.ops import nms, batched_nms, box_iou

# Two scooters parked close together. The detector fired five times.
boxes = torch.tensor([
    [100., 100., 200., 260.],   # scooter A, best guess
    [108., 106., 194., 258.],   # A again
    [ 90.,  94., 214., 268.],   # A again, looser
    [140., 100., 240., 260.],   # scooter B, best guess
    [147., 104., 235., 256.],   # B again
])
scores = torch.tensor([0.92, 0.88, 0.75, 0.81, 0.79])

def nms_by_hand(boxes, scores, thresh):
    order = scores.argsort(descending=True)          # loudest claim goes first
    keep = []
    while order.numel() > 0:
        i = order[0].item()
        keep.append(i)
        rest = order[1:]
        if rest.numel() == 0:
            break
        ious = box_iou(boxes[i:i + 1], boxes[rest]).squeeze(0)
        order = rest[ious <= thresh]                 # anything hugging the winner is deleted
    return torch.tensor(keep)

print("kept by hand       :", nms_by_hand(boxes, scores, 0.5).tolist())
print("kept by torchvision:", nms(boxes, scores, 0.5).tolist())

print("\nthe two boxes that survive at 0.5:")
for i in nms(boxes, scores, 0.5).tolist():
    print(f"  box {i}  score {scores[i]:.2f}  {[int(v) for v in boxes[i].tolist()]}")

print("\nthe same five boxes at different thresholds:")
for t in [0.3, 0.4, 0.5, 0.7, 0.9]:
    print(f"  iou_threshold={t}  ->  kept {nms(boxes, scores, t).tolist()}")

print("\nhow much each box overlaps scooter A's best box:")
print("  ", [round(v, 3) for v in box_iou(boxes[:1], boxes).squeeze(0).tolist()])

# Different classes stacked on top of each other.
stacked = torch.tensor([[10., 10., 110., 210.],      # a person
                        [12., 60., 108., 200.]])     # the bag they carry
ssc = torch.tensor([0.90, 0.85])
print(f"\nperson and bag overlap at IoU {box_iou(stacked[:1], stacked[1:]).item():.2f}")
print("  class-blind nms keeps:", nms(stacked, ssc, 0.5).tolist(), " <- the bag is deleted")
print("  batched_nms keeps    :",
      sorted(batched_nms(stacked, ssc, torch.tensor([0, 26]), 0.5).tolist()))

# Soft-NMS: lower the neighbour's score instead of deleting it.
def soft_nms(boxes, scores, sigma=0.5, min_score=0.05):
    s, out = scores.clone(), []
    while (s > min_score).any():
        i = torch.where(s > min_score, s, torch.full_like(s, -1.0)).argmax()
        out.append((i.item(), round(s[i].item(), 3)))
        ious = box_iou(boxes[i:i + 1], boxes).squeeze(0)
        s = s * torch.exp(-(ious ** 2) / sigma)      # gaussian penalty, nothing is ever deleted
        s[i] = 0.0
    return out

print("\nsoft-nms, index and surviving score:")
for i, sc in soft_nms(boxes, scores):
    print(f"   box {i}  {sc}")

import torch
from torchvision.ops import batched_nms, clip_boxes_to_image, remove_small_boxes

torch.manual_seed(0)
N = 20_000
boxes = torch.rand(N, 4) * 640
boxes[:, 2:] += boxes[:, :2]                       # make x2 > x1 and y2 > y1
scores = torch.rand(N) ** 8                        # most scores near zero, a few high
labels = torch.randint(0, 80, (N,))

print("raw candidates                :", len(boxes))

keep = scores > 0.05                               # 1. score threshold, does most of the work
boxes, scores, labels = boxes[keep], scores[keep], labels[keep]
print("after score threshold 0.05    :", len(boxes))

order = scores.argsort(descending=True)[:1000]     # 2. keep only the top 1000 into NMS
boxes, scores, labels = boxes[order], scores[order], labels[order]
print("after top-k truncation        :", len(boxes))

boxes = clip_boxes_to_image(boxes, (640, 640))     # 3. tidy the geometry
keep = remove_small_boxes(boxes, min_size=2)
boxes, scores, labels = boxes[keep], scores[keep], labels[keep]
print("after clipping and size filter:", len(boxes))

keep = batched_nms(boxes, scores, labels, 0.5)[:100]   # 4. NMS, then cap the output
print("after per-class NMS, capped   :", len(keep))