megadetector 10.0.13__py3-none-any.whl

megadetector/tests/test_nms_synthetic.py

@@ -0,0 +1,335 @@
+"""
+
+Test script for validating NMS functionality with synthetic data.
+
+This script creates synthetic detection scenarios where we know exactly which
+boxes should be suppressed by NMS, allowing us to verify the correctness of
+the NMS implementation.
+
+This is an AI-generated test module.
+
+"""
+
+
+#%% Imports
+
+import torch
+
+from megadetector.detection.pytorch_detector import nms
+
+
+#%% Support functions
+
+def calculate_iou_boxes(box1, box2):
+    """
+    Calculate IoU between two boxes in [x1, y1, x2, y2] format.
+
+    Args:
+        box1: torch.Tensor or list of [x1, y1, x2, y2]
+        box2: torch.Tensor or list of [x1, y1, x2, y2]
+
+    Returns:
+        float: IoU value between 0 and 1
+    """
+
+    if isinstance(box1, list):
+        box1 = torch.tensor(box1, dtype=torch.float)
+    if isinstance(box2, list):
+        box2 = torch.tensor(box2, dtype=torch.float)
+
+    # Calculate intersection area
+    x1_inter = max(box1[0], box2[0])
+    y1_inter = max(box1[1], box2[1])
+    x2_inter = min(box1[2], box2[2])
+    y2_inter = min(box1[3], box2[3])
+
+    if x2_inter <= x1_inter or y2_inter <= y1_inter:
+        return 0.0
+
+    intersection = (x2_inter - x1_inter) * (y2_inter - y1_inter)
+
+    # Calculate union area
+    area1 = (box1[2] - box1[0]) * (box1[3] - box1[1])
+    area2 = (box2[2] - box2[0]) * (box2[3] - box2[1])
+    union = area1 + area2 - intersection
+
+    return float(intersection / union) if union > 0 else 0.0
+
+
+def create_synthetic_predictions():
+    """
+    Create synthetic model predictions for testing NMS.
+
+    Returns:
+        torch.Tensor: Synthetic predictions in the format expected by the NMS function
+                     Shape: [batch_size=1, num_anchors, num_classes + 5]
+
+    Test scenarios:
+    1. Two highly overlapping boxes (IoU > 0.5) with different confidences - higher confidence should win
+    2. Two boxes with low overlap (IoU < 0.5) - both should be kept
+    3. Multiple boxes of different classes in same location - should be kept (class-independent NMS)
+    4. Three overlapping boxes with cascading confidences - highest confidence should win
+    """
+
+    # We'll create predictions for a 640x640 image with 3 classes
+    # Format: [x_center, y_center, width, height, objectness, class0_conf, class1_conf, class2_conf]
+
+    synthetic_boxes = []
+
+    # Scenario 1: Two highly overlapping boxes (IoU > 0.8)
+    # Box A: center=(100, 100), size=80x80, high confidence for class 0
+    # Box B: center=(105, 105), size=80x80, low confidence for class 0  (smaller offset = higher IoU)
+    # Expected: Box A kept, Box B suppressed
+    synthetic_boxes.append([100, 100, 80, 80, 0.9, 0.8, 0.1, 0.1])  # Box A - should be kept
+    synthetic_boxes.append([105, 105, 80, 80, 0.9, 0.5, 0.1, 0.1])  # Box B - should be suppressed
+
+    # Scenario 1b: Two nearly identical boxes (IoU ≈ 0.95)
+    # Box A2: center=(200, 100), size=60x60, high confidence for class 0
+    # Box B2: center=(202, 102), size=60x60, lower confidence for class 0
+    # Expected: Box A2 kept, Box B2 suppressed
+    synthetic_boxes.append([200, 100, 60, 60, 0.9, 0.9, 0.05, 0.05])  # Box A2 - should be kept
+    synthetic_boxes.append([202, 102, 60, 60, 0.9, 0.7, 0.1, 0.1])    # Box B2 - should be suppressed
+
+    # Scenario 2: Two boxes with low overlap (IoU ≈ 0.1)
+    # Box C: center=(300, 100), size=60x60, class 0
+    # Box D: center=(380, 100), size=60x60, class 0
+    # Expected: Both kept
+    synthetic_boxes.append([300, 100, 60, 60, 0.9, 0.7, 0.1, 0.1])  # Box C - should be kept
+    synthetic_boxes.append([380, 100, 60, 60, 0.9, 0.6, 0.1, 0.1])  # Box D - should be kept
+
+    # Scenario 3: Same location, different classes
+    # Box E: center=(100, 300), size=70x70, class 0
+    # Box F: center=(100, 300), size=70x70, class 1
+    # Expected: Both kept (class-independent NMS)
+    synthetic_boxes.append([100, 300, 70, 70, 0.9, 0.7, 0.1, 0.1])  # Box E - class 0, should be kept
+    synthetic_boxes.append([100, 300, 70, 70, 0.9, 0.1, 0.7, 0.1])  # Box F - class 1, should be kept
+
+    # Scenario 4: Three cascading overlapping boxes
+    # Box G: center=(500, 300), size=80x80, highest confidence
+    # Box H: center=(510, 310), size=80x80, medium confidence
+    # Box I: center=(520, 320), size=80x80, lowest confidence
+    # Expected: Only Box G kept
+    synthetic_boxes.append([500, 300, 80, 80, 0.95, 0.9, 0.05, 0.05])  # Box G - highest conf, should be kept
+    synthetic_boxes.append([510, 310, 80, 80, 0.9,  0.7, 0.1,  0.1])   # Box H - should be suppressed
+    synthetic_boxes.append([520, 320, 80, 80, 0.85, 0.6, 0.15, 0.15])  # Box I - should be suppressed
+
+    # Add some low-confidence boxes that should be filtered out before NMS
+    synthetic_boxes.append([200, 500, 50, 50, 0.1, 0.05, 0.02, 0.03])  # Too low confidence
+
+    # Convert to tensor format expected by NMS function
+    # We need to pad to a reasonable number of anchors (let's use 20)
+    num_anchors = 20
+    num_classes = 3
+
+    predictions = torch.zeros(1, num_anchors, num_classes + 5)  # batch_size=1
+
+    # Fill in our synthetic boxes
+    for i, box_data in enumerate(synthetic_boxes):
+        if i < num_anchors:
+            predictions[0, i, :] = torch.tensor(box_data)
+
+    return predictions
+
+
+#%% Main test function
+
+def test_nms_functionality():
+    """
+    Test the NMS function with synthetic data to verify correct behavior.
+    """
+
+    print("Testing NMS functionality with synthetic data...")
+
+    # Generate synthetic predictions
+    predictions = create_synthetic_predictions()
+    print(f"Created synthetic predictions with shape: {predictions.shape}")
+
+    # Run NMS with IoU threshold = 0.5 and confidence threshold = 0.3
+    results = nms(predictions, conf_thres=0.3, iou_thres=0.5, max_det=300)
+
+    print(f"NMS returned {len(results)} batch results")
+    detections = results[0]  # Get results for first (and only) image
+    print(f"Number of detections after NMS: {detections.shape[0]}")
+
+    assert detections.shape[0] != 0
+
+    print("\nDetections after NMS:")
+    print("Format: [x1, y1, x2, y2, confidence, class_id]")
+    for i, det in enumerate(detections):
+        x1, y1, x2, y2, conf, cls = det
+        center_x = (x1 + x2) / 2
+        center_y = (y1 + y2) / 2
+        width = x2 - x1
+        height = y2 - y1
+        print(f"Detection {i}: center=({center_x:.1f}, {center_y:.1f}), "
+              f"size={width:.1f}x{height:.1f}, conf={conf:.3f}, class={int(cls)}")
+
+    # Verify expected results
+
+    # Verify that high-confidence boxes are kept over low-confidence overlapping ones
+    # Look for the scenario 1 boxes (around center 100,100 area)
+    scenario1_boxes = []
+    for i, det in enumerate(detections):
+        x1, y1, x2, y2, conf, cls = det
+        center_x = (x1 + x2) / 2
+        center_y = (y1 + y2) / 2
+        if 80 <= center_x <= 130 and 80 <= center_y <= 130 and int(cls) == 0:
+            scenario1_boxes.append((i, center_x, center_y, conf))
+
+    # Check scenario 1b (around center 200,100 area)
+    scenario1b_boxes = []
+    for i, det in enumerate(detections):
+        x1, y1, x2, y2, conf, cls = det
+        center_x = (x1 + x2) / 2
+        center_y = (y1 + y2) / 2
+        if 180 <= center_x <= 220 and 80 <= center_y <= 120 and int(cls) == 0:
+            scenario1b_boxes.append((i, center_x, center_y, conf))
+
+    # Both scenario 1 and 1b should have exactly 1 detection each
+    total_high_overlap_boxes = len(scenario1_boxes) + len(scenario1b_boxes)
+    if total_high_overlap_boxes != 2:
+        print("Error: expected 2 detections in high-overlap scenarios (1 each), got {}".format(
+            total_high_overlap_boxes
+        ))
+        print(f"  Scenario 1: {len(scenario1_boxes)} boxes")
+        print(f"  Scenario 1b: {len(scenario1b_boxes)} boxes")
+        raise AssertionError()
+    # Should be the high-confidence box (0.8 * 0.9 = 0.72)
+    elif len(scenario1_boxes) == 1 and scenario1_boxes[0][3] < 0.7:
+        print("Error: wrong box kept in scenario 1. Expected conf > 0.7, got {}".format(
+            scenario1_boxes[0][3]
+        ))
+        raise AssertionError()
+    # Should be the high-confidence box (0.9 * 0.9 = 0.81)
+    elif len(scenario1b_boxes) == 1 and scenario1b_boxes[0][3] < 0.8:
+        print("Error: wrong box kept in scenario 1b. Expected conf > 0.8, got {}".format(
+            scenario1b_boxes[0][3]
+        ))
+        raise AssertionError()
+    else:
+        print("Scenarios 1 & 1b passed: High-confidence boxes kept, low-confidence overlapping boxes suppressed")
+
+        # Verify IoU calculations and ensure suppression actually works
+        if len(scenario1_boxes) == 1 and len(scenario1b_boxes) == 1:
+            # Calculate what the IoU would have been between the boxes that were supposed to overlap
+            # Scenario 1: Box A (100,100,80x80) vs Box B (105,105,80x80)
+            box_a = [100-40, 100-40, 100+40, 100+40]  # Convert center+size to corners
+            box_b = [105-40, 105-40, 105+40, 105+40]
+            iou_1 = calculate_iou_boxes(box_a, box_b)
+
+            # Scenario 1b: Box A2 (200,100,60x60) vs Box B2 (202,102,60x60)
+            box_a2 = [200-30, 100-30, 200+30, 100+30]
+            box_b2 = [202-30, 102-30, 202+30, 102+30]
+            iou_1b = calculate_iou_boxes(box_a2, box_b2)
+
+            print(f"    Theoretical IoU for scenario 1 boxes: {iou_1:.3f}")
+            print(f"    Theoretical IoU for scenario 1b boxes: {iou_1b:.3f}")
+
+            # If IoU > threshold, suppression should have occurred
+            if iou_1 <= 0.5:
+                print(f"Error: scenario 1 IoU {iou_1:.3f} is too low - test setup is invalid!")
+                raise AssertionError()
+            elif iou_1b <= 0.5:
+                print(f"Error: scenario 1b IoU {iou_1b:.3f} is too low - test setup is invalid!")
+                raise AssertionError()
+            else:
+                print("    High IoU confirmed - suppression was correct")
+
+    # Verify scenario 2 - both non-overlapping boxes should be kept
+    scenario2_boxes = []
+    for i, det in enumerate(detections):
+        x1, y1, x2, y2, conf, cls = det
+        center_x = (x1 + x2) / 2
+        center_y = (y1 + y2) / 2
+        if 270 <= center_x <= 410 and 70 <= center_y <= 130 and int(cls) == 0:
+            scenario2_boxes.append((i, center_x, center_y, conf))
+
+    if len(scenario2_boxes) != 2:
+        print(f"Error: expected 2 detections in scenario 2 area, got {len(scenario2_boxes)}")
+        raise AssertionError()
+    else:
+        print("Scenario 2 passed: Both non-overlapping boxes kept")
+
+    # Verify scenario 3 - different classes should both be kept
+    scenario3_boxes = []
+    for i, det in enumerate(detections):
+        x1, y1, x2, y2, conf, cls = det
+        center_x = (x1 + x2) / 2
+        center_y = (y1 + y2) / 2
+        if 65 <= center_x <= 135 and 265 <= center_y <= 335:
+            scenario3_boxes.append((i, center_x, center_y, conf, int(cls)))
+
+    classes_found = set(box[4] for box in scenario3_boxes)
+    if (len(scenario3_boxes) != 2) or (len(classes_found) != 2):
+        print("Error: expected 2 detections of different classes , got {} detections of classes {}".format(
+            len(scenario3_boxes),classes_found
+        ))
+        raise AssertionError()
+    else:
+        print("Scenario 3 passed: Both different-class boxes kept")
+
+    # Verify scenario 4 - cascading overlapping boxes (only highest confidence should remain)
+    scenario4_boxes = []
+    for i, det in enumerate(detections):
+        x1, y1, x2, y2, conf, cls = det
+        center_x = (x1 + x2) / 2
+        center_y = (y1 + y2) / 2
+        if 460 <= center_x <= 560 and 260 <= center_y <= 360 and int(cls) == 0:
+            scenario4_boxes.append((i, center_x, center_y, conf))
+
+    print(f"\nScenario 4 analysis: Found {len(scenario4_boxes)} boxes in cascading area:")
+    for i, (det_idx, cx, cy, conf) in enumerate(scenario4_boxes):
+        print(f"  Box {i}: center=({cx:.1f}, {cy:.1f}), conf={conf:.3f}")
+
+    # Check IoU between remaining boxes to ensure proper suppression
+    if len(scenario4_boxes) >= 2:
+        max_iou = 0
+        for i in range(len(scenario4_boxes)):
+            for j in range(i+1, len(scenario4_boxes)):
+                det1 = detections[scenario4_boxes[i][0]]
+                det2 = detections[scenario4_boxes[j][0]]
+                iou = calculate_iou_boxes(det1[:4], det2[:4])
+                print(f"  IoU between box {i} and box {j}: {iou:.3f}")
+                max_iou = max(max_iou, iou)
+
+        if len(scenario4_boxes) == 1:
+            print("Scenario 4 passed: Only highest confidence box kept")
+        else:
+            # This is only OK if IoU < threshold
+            if max_iou < 0.5:  # Our IoU threshold
+                print("Scenario 4 passed: Multiple boxes kept due to low IoU (< 0.5)")
+            else:
+                print(f"ERROR: Scenario 4 failed - boxes with IoU {max_iou:.3f} > 0.5 were not suppressed!")
+                raise AssertionError()
+
+    # Create a scenario that requires IoU calculation
+    print("\n=== COMPREHENSIVE IoU VALIDATION TEST ===")
+
+    # Create two identical boxes that should definitely be suppressed
+    identical_box_a = [100, 100, 50, 50, 0.9, 0.9, 0.05, 0.05]  # High confidence
+    identical_box_b = [100, 100, 50, 50, 0.9, 0.7, 0.1, 0.1]    # Lower confidence
+
+    test_predictions = torch.zeros(1, 5, 8)  # Small batch for focused test
+    test_predictions[0, 0, :] = torch.tensor(identical_box_a)
+    test_predictions[0, 1, :] = torch.tensor(identical_box_b)
+
+    # Run NMS on this simple case
+    test_results = nms(test_predictions, conf_thres=0.3, iou_thres=0.5, max_det=300)
+    test_detections = test_results[0]
+
+    print(f"Identical boxes test: Input 2 identical boxes, got {test_detections.shape[0]} detections")
+
+    if test_detections.shape[0] != 1:
+        print(f"Error Two identical boxes should result in 1 detection, got {test_detections.shape[0]}")
+        raise AssertionError()
+    else:
+        # Verify it kept the higher confidence box
+        kept_conf = test_detections[0, 4].item()
+        expected_conf = 0.9 * 0.9  # objectness * class_conf
+        if abs(kept_conf - expected_conf) > 0.01:
+            print(f"ERROR: Wrong box kept. Expected conf ≈ {expected_conf:.3f}, got {kept_conf:.3f}")
+            raise AssertionError()
+        else:
+            print("Identical boxes test passed: Higher confidence box kept")
+
+    print("\nNMS tests passed")