add wbf

you immediately see what subject materials im reading and where im at

pipeline/ops/__init__.py

@@ -17,12 +17,14 @@ from pipeline.ops.geometry import (
     point_in_box,
     point_in_polygon,
 )
-from pipeline.ops.nms import nms, nms_by_class
+from pipeline.ops.nms import nms, nms_by_class, wbf, wbf_by_class
 
 __all__ = [
-    # NMS
+    # NMS / Box Fusion
     "nms",
     "nms_by_class",
+    "wbf",
+    "wbf_by_class",
     # Geometry
     "expand_box",
     "point_in_box",

pipeline/ops/nms.py

@@ -1,4 +1,4 @@
-"""Non-Maximum Suppression algorithms for detection filtering."""
+"""Box filtering algorithms: NMS and Weighted Boxes Fusion."""
 
 from typing import List
 
@@ -59,3 +59,119 @@ def nms_by_class(
         result.extend(nms(dets, iou_threshold))
 
     return result
+
+
+def wbf(
+    detections: List[Detection],
+    iou_threshold: float = 0.5,
+    skip_threshold: float = 0.0,
+) -> List[Detection]:
+    """
+    Weighted Boxes Fusion: Merge overlapping boxes instead of suppressing.
+
+    Unlike NMS which picks the best box and discards overlapping ones,
+    WBF fuses overlapping boxes by averaging their coordinates weighted
+    by confidence. This often gives better localization.
+
+    Args:
+        detections: List of Detection objects
+        iou_threshold: Boxes with IoU > threshold are fused together
+        skip_threshold: Ignore boxes with confidence below this
+
+    Returns:
+        Fused list of detections
+
+    Reference:
+        https://arxiv.org/abs/1910.13302
+    """
+    if not detections:
+        return []
+
+    # Filter low confidence
+    dets = [d for d in detections if d.score >= skip_threshold]
+    if not dets:
+        return []
+
+    # Sort by confidence (highest first)
+    dets = sorted(dets, key=lambda d: d.score, reverse=True)
+
+    # Clusters of fused boxes
+    clusters: List[List[Detection]] = []
+
+    for det in dets:
+        matched_cluster = None
+
+        # Find a cluster this box belongs to
+        for cluster in clusters:
+            # Check IoU with the fused box (first element is always the fused result)
+            if cluster[0].box.iou(det.box) >= iou_threshold:
+                matched_cluster = cluster
+                break
+
+        if matched_cluster is not None:
+            # Add to existing cluster and recompute fused box
+            matched_cluster.append(det)
+            _update_fused_box(matched_cluster)
+        else:
+            # Start new cluster
+            clusters.append([det])
+
+    # Return the fused box from each cluster (first element)
+    return [cluster[0] for cluster in clusters]
+
+
+def _update_fused_box(cluster: List[Detection]):
+    """
+    Recompute the fused box for a cluster.
+    The fused box is stored as cluster[0].
+    """
+    if len(cluster) == 1:
+        return
+
+    # Weighted average of coordinates
+    total_weight = sum(d.score for d in cluster)
+
+    x1 = sum(d.box.x1 * d.score for d in cluster) / total_weight
+    y1 = sum(d.box.y1 * d.score for d in cluster) / total_weight
+    x2 = sum(d.box.x2 * d.score for d in cluster) / total_weight
+    y2 = sum(d.box.y2 * d.score for d in cluster) / total_weight
+
+    # Fused confidence: average boosted by number of boxes
+    # More boxes agreeing = higher confidence (capped at 1.0)
+    avg_conf = total_weight / len(cluster)
+    fused_conf = min(1.0, avg_conf * (1 + 0.1 * (len(cluster) - 1)))
+
+    # Update the first element with fused values
+    fused = cluster[0]
+    fused.box.x1 = int(x1)
+    fused.box.y1 = int(y1)
+    fused.box.x2 = int(x2)
+    fused.box.y2 = int(y2)
+    fused.box.conf = fused_conf
+    fused.score = fused_conf
+    fused.metadata["wbf_count"] = len(cluster)
+
+
+def wbf_by_class(
+    detections: List[Detection],
+    iou_threshold: float = 0.5,
+    skip_threshold: float = 0.0,
+) -> List[Detection]:
+    """
+    Run WBF separately for each class label.
+    This prevents boxes of different classes from being fused together.
+    """
+    if not detections:
+        return []
+
+    # Group by label
+    by_label: dict[str, List[Detection]] = {}
+    for det in detections:
+        by_label.setdefault(det.label, []).append(det)
+
+    # Run WBF per class
+    result = []
+    for dets in by_label.values():
+        result.extend(wbf(dets, iou_threshold, skip_threshold))
+
+    return result

pipeline/stages/api.py

@@ -148,7 +148,7 @@ class APIPersonDetector(APIStageBase):
         resp.raise_for_status()
         return resp.json()
 
-
+# haiku 4.5 is multimodal but i fear its going to ping 65 million images per second (we need cache control? who are we? why cant we upload video?)
 class VLMVerifier(APIStageBase):
     """
     Vision-Language Model verification stage.

pipeline/stages/runner.py

@@ -60,7 +60,7 @@ class StagePipeline:
                     stage.name,
                     ctx.stop_reason,
                 )
-                continue
+                continue # break?
 
             try:
                 result = stage.process(ctx)