5 · NMS 的阈值不是评价的阈值 · 实现代码

def class_aware_nms(boxes, scores, classes, iou_threshold=0.50):
    """Keep a box unless a higher-scored kept box of SAME class has IoU > threshold.

    Equality is retained. There is no extra confidence cutoff in this lab.
    """
    boxes = boxes_xyxy(boxes)
    scores, classes = np.asarray(scores, float), np.asarray(classes)
    if scores.shape != (len(boxes),) or classes.shape != (len(boxes),):
        raise ValueError("Scores/classes must have shape (N,)")
    if not np.isfinite(scores).all() or not 0 <= iou_threshold <= 1:
        raise ValueError("Finite scores and an IoU threshold in [0,1] required")
    remaining = np.argsort(-scores, kind="stable")
    kept, suppressed = [], []
    while len(remaining):
        current, remaining = int(remaining[0]), remaining[1:]
        kept.append(current)
        overlap = box_iou(boxes[current:current + 1], boxes[remaining])[0]
        remove = (classes[remaining] == classes[current]) & (overlap > iou_threshold)
        for victim, iou_value in zip(remaining[remove], overlap[remove]):
            suppressed.append({"prediction": int(victim + 1), "by": current + 1,
                               "iou": float(iou_value), "class": int(classes[current])})
        remaining = remaining[~remove]
    return np.asarray(kept, dtype=int), suppressed

print("NMS keeps 1-based predictions:", class_aware_nms(PRED, SCORES, CLASSES, NMS_THRESHOLD)[0] + 1)