Non-maximum suppression in PyTorch - ML Experiment: Train & Evaluate

Experiment - Non-maximum suppression

Problem:You have a set of bounding boxes predicted by an object detection model. Many boxes overlap heavily, causing multiple detections for the same object.

Current Metrics:Precision: 75%, Recall: 80%, with many duplicate detections due to overlapping boxes.

Issue:The model outputs multiple overlapping boxes for the same object, reducing precision and making results confusing.

Your Task

Implement Non-maximum suppression (NMS) to remove overlapping bounding boxes and improve precision to at least 85% while maintaining recall above 75%.

Use PyTorch for implementation.

Do not change the model architecture or training data.

Only modify the post-processing step to include NMS.

Hint 1

Hint 2

Hint 3

Hint 4

Solution

PyTorch

import torch

def non_max_suppression(boxes, scores, iou_threshold=0.5):
    # boxes: tensor of shape (N, 4) with (x1, y1, x2, y2)
    # scores: tensor of shape (N,) with confidence scores
    # returns indices of boxes to keep

    # Sort boxes by scores descending
    sorted_indices = torch.argsort(scores, descending=True)
    keep = []

    while sorted_indices.numel() > 0:
        # Pick the box with highest score
        current = sorted_indices[0]
        keep.append(current.item())

        if sorted_indices.numel() == 1:
            break

        current_box = boxes[current].unsqueeze(0)  # shape (1,4)
        other_boxes = boxes[sorted_indices[1:]]  # shape (N-1,4)

        # Compute IoU
        x1 = torch.max(current_box[:, 0], other_boxes[:, 0])
        y1 = torch.max(current_box[:, 1], other_boxes[:, 1])
        x2 = torch.min(current_box[:, 2], other_boxes[:, 2])
        y2 = torch.min(current_box[:, 3], other_boxes[:, 3])

        inter_w = (x2 - x1).clamp(min=0)
        inter_h = (y2 - y1).clamp(min=0)
        inter_area = inter_w * inter_h

        area_current = (current_box[:, 2] - current_box[:, 0]) * (current_box[:, 3] - current_box[:, 1])
        area_others = (other_boxes[:, 2] - other_boxes[:, 0]) * (other_boxes[:, 3] - other_boxes[:, 1])

        union_area = area_current + area_others - inter_area
        iou = inter_area / union_area

        # Keep boxes with IoU less than threshold
        below_threshold = iou <= iou_threshold
        sorted_indices = sorted_indices[1:][below_threshold]

    return keep

# Example usage:
boxes = torch.tensor([
    [10, 10, 50, 50],
    [12, 12, 48, 48],
    [100, 100, 150, 150],
    [110, 110, 140, 140]
], dtype=torch.float32)

scores = torch.tensor([0.9, 0.85, 0.8, 0.7])

keep_indices = non_max_suppression(boxes, scores, iou_threshold=0.5)
print("Boxes kept after NMS:", keep_indices)

Added a post-processing step using Non-maximum suppression to remove overlapping bounding boxes.

Implemented IoU calculation to measure overlap between boxes.

Filtered out boxes with IoU above 0.5 to keep only the most confident detections.

Results Interpretation

Before NMS: Precision 75%, Recall 80%, many overlapping boxes causing duplicates.

After NMS: Precision improved to 87%, Recall slightly decreased to 78%, duplicates removed.

Non-maximum suppression helps clean up overlapping predictions by keeping only the best box per object, improving precision and making detection results clearer.

Bonus Experiment

Try adjusting the IoU threshold from 0.3 to 0.7 and observe how precision and recall change.

💡 Hint

Lower IoU threshold removes more boxes (higher precision, lower recall), higher threshold keeps more boxes (higher recall, lower precision). Find the best balance.

Practice

1. What is the main purpose of Non-maximum Suppression (NMS) in object detection using PyTorch?

easy

A. To increase the number of detected boxes

B. To resize images before detection

C. To randomly select boxes for training

D. To remove overlapping bounding boxes and keep only the best ones

Solution

1. Step 1: Understand the role of NMS

   Non-maximum suppression is used to filter out overlapping bounding boxes that represent the same object.

2. Step 2: Identify the goal of NMS

   It keeps only the bounding box with the highest confidence score among overlapping boxes to avoid duplicate detections.

3. Final Answer:

   To remove overlapping bounding boxes and keep only the best ones -> Option D

4. Quick Check:

   NMS purpose = remove overlaps, keep best [OK]

Hint: NMS filters duplicates, keeping only top scoring boxes [OK]

Common Mistakes:

• Thinking NMS increases detected boxes
• Confusing NMS with image resizing
• Assuming NMS randomly selects boxes

2. Which PyTorch function is commonly used to perform Non-maximum Suppression on bounding boxes and scores?

easy

A. torch.sort

B. torch.max

C. torch.ops.torchvision.nms

D. torch.argmax

Solution

1. Step 1: Identify the PyTorch NMS function

   PyTorch provides a dedicated function called torch.ops.torchvision.nms to perform Non-maximum Suppression.

2. Step 2: Differentiate from other functions

   torch.max, torch.sort, and torch.argmax do not perform NMS; they serve different purposes like finding max values or sorting tensors.

3. Final Answer:

   torch.ops.torchvision.nms -> Option C

4. Quick Check:

   PyTorch NMS function = torch.ops.torchvision.nms [OK]

Hint: Use torch.ops.torchvision.nms for Non-maximum Suppression in PyTorch [OK]

Common Mistakes:

• Using torch.max instead of torch.ops.torchvision.nms
• Confusing sorting functions with NMS
• Trying torch.argmax for NMS

3. Given the following PyTorch code snippet, what will be the output tensor of indices after applying NMS?

import torch
boxes = torch.tensor([[10, 10, 20, 20], [11, 11, 21, 21], [30, 30, 40, 40]], dtype=torch.float)
scores = torch.tensor([0.9, 0.8, 0.7])
iou_threshold = 0.5
indices = torch.ops.torchvision.nms(boxes, scores, iou_threshold)
print(indices.tolist())

medium

A. [0, 2]

B. [1, 2]

C. [0, 1, 2]

D. [2]

Solution

1. Step 1: Analyze box overlaps and scores

   Boxes 0 and 1 overlap with IoU > 0.5, box 2 is separate. Box 0 has higher score (0.9) than box 1 (0.8).

2. Step 2: Apply NMS logic

   NMS keeps box 0, removes box 1 due to overlap and lower score, and keeps box 2 as it does not overlap.

3. Final Answer:

   [0, 2] -> Option A

4. Quick Check:

   Keep highest score boxes without high overlap = [0, 2] [OK]

Hint: Keep highest score boxes, remove overlapping lower scores [OK]

Common Mistakes:

• Keeping all boxes ignoring overlap
• Removing box 0 instead of box 1
• Keeping only one box

4. Identify the error in this PyTorch NMS usage and select the correct fix:

import torch
boxes = torch.tensor([[0, 0, 10, 10], [1, 1, 11, 11]], dtype=torch.float)
scores = torch.tensor([0.5, 0.6])
iou_threshold = 0.3
indices = torch.nms(boxes, scores, iou_threshold)
print(indices)

medium

A. Set iou_threshold to 1.0

B. Use torch.ops.torchvision.nms instead of torch.nms

C. Change boxes dtype to int

D. Sort scores before calling nms

Solution

1. Step 1: Check function availability

   torch.nms is not a direct PyTorch function; NMS is accessed via torchvision ops as torch.ops.torchvision.nms.

2. Step 2: Correct function call

   Replacing torch.nms with torch.ops.torchvision.nms fixes the error and runs NMS correctly.

3. Final Answer:

   Use torch.ops.torchvision.nms instead of torch.nms -> Option B

4. Quick Check:

   Correct NMS call = torch.ops.torchvision.nms [OK]

Hint: Use torchvision's nms via torch.ops.torchvision.nms [OK]

Common Mistakes:

• Calling torch.nms directly
• Changing box dtype unnecessarily
• Adjusting IoU threshold without fixing function

5. You have detected multiple overlapping boxes with scores: [0.95, 0.9, 0.85, 0.6]. If you apply Non-maximum Suppression with a very low IoU threshold (e.g., 0.1), what is the expected effect on the number of boxes kept?

hard

A. Most boxes will be removed because even slight overlaps are suppressed

B. All boxes will be kept because the threshold is low

C. Only the box with the lowest score will be kept

D. Boxes will be randomly removed regardless of overlap

Solution

1. Step 1: Understand IoU threshold meaning

   A low IoU threshold means boxes with even small overlaps are considered duplicates and suppressed.

2. Step 2: Predict effect on boxes

   Because boxes overlap slightly, most will be removed except the highest scoring box, reducing the number kept.

3. Final Answer:

   Most boxes will be removed because even slight overlaps are suppressed -> Option A

4. Quick Check:

   Low IoU threshold = strict filtering = fewer boxes [OK]

Hint: Lower IoU threshold means stricter overlap removal [OK]

Common Mistakes:

• Thinking low threshold keeps more boxes
• Assuming random removal
• Believing lowest score box is kept