Postprocessing Backends¶

def set_postprocess_backend(name: str) -> None:
    """Set the postprocessing backend.

    Call once at startup before running any inference.  This function is
    **not** thread-safe.

    Args:
        name: One of "auto", "numpy", "numba", "torchvision".
    """
    global _backend, _resolved_cache
    if name not in VALID_BACKENDS:
        raise ValueError(f"Unknown backend {name!r}. Choose from {VALID_BACKENDS}")
    _backend = name
    _resolved_cache = None  # force re-resolve
    # Invalidate the dispatch function cache in combine module
    try:
        from sahi.postprocess.combine import _dispatch_cache

        _dispatch_cache.clear()
    except ImportError:
        pass

def get_postprocess_backend() -> str:
    """Return the currently configured backend name (may be "auto")."""
    return _backend

def resolve_backend() -> str:
    """Resolve "auto" to a concrete backend, caching the result.

    When the backend is set to "auto", detection follows this priority:

    1. **torchvision** -- selected if torchvision is installed and a GPU is
       available, either CUDA or Apple MPS (GPU-accelerated NMS).
    2. **numba** -- selected if the numba package is installed (JIT-compiled
       loops, faster than pure numpy for large prediction counts).
    3. **numpy** -- always available as the fallback (pure numpy,
       no extra dependencies).

    If the backend was explicitly set via ``set_postprocess_backend``, that
    value is returned directly without auto-detection.

    Returns:
        One of "numpy", "numba", or "torchvision".
    """
    global _resolved_cache
    if _resolved_cache is not None:
        return _resolved_cache

    if _backend != "auto":
        _resolved_cache = _backend
        return _backend

    # Auto-detect: prefer torchvision on CUDA/MPS, then numba, then numpy
    if is_available("torchvision"):
        try:
            import torch

            # torch.mps.is_available() only exists from torch 2.5; torch.backends.mps
            # has been the stable entry point since 1.12 and is False on non-Apple builds.
            if torch.cuda.is_available() or torch.backends.mps.is_available():
                _resolved_cache = "torchvision"
                return _resolved_cache
        except ImportError:
            pass

    if is_available("numba"):
        _resolved_cache = "numba"
        return _resolved_cache

    _resolved_cache = "numpy"
    return _resolved_cache

def nms(
    predictions: np.ndarray,
    match_metric: str = "IOU",
    match_threshold: float = 0.5,
) -> list[int]:
    """Non-maximum suppression for axis-aligned bounding boxes.

    Dispatches to the resolved backend (numpy, numba, or torchvision).

    Args:
        predictions: Array of shape (N, 6) with columns
            [x1, y1, x2, y2, score, category_id].
        match_metric: Overlap metric, "IOU" or "IOS".
        match_threshold: Minimum overlap to suppress a lower-scored box.

    Returns:
        List of indices of the kept predictions, sorted by score descending.
    """
    if len(predictions) == 0:
        return []
    return _dispatch("nms")(predictions, match_metric, match_threshold)

def batched_nms(
    predictions: np.ndarray,
    match_metric: str = "IOU",
    match_threshold: float = 0.5,
) -> list[int]:
    """Apply non-maximum suppression independently per category.

    Args:
        predictions: Array of shape (N, 6) with columns
            [x1, y1, x2, y2, score, category_id].
        match_metric: Overlap metric, "IOU" or "IOS".
        match_threshold: Minimum overlap to suppress a lower-scored box.

    Returns:
        List of indices of the kept predictions, sorted by score descending.
    """
    return _batched_apply(predictions, nms, match_metric, match_threshold)  # type: ignore[return-value]

def greedy_nmm(
    predictions: np.ndarray,
    match_metric: str = "IOU",
    match_threshold: float = 0.5,
) -> dict[int, list[int]]:
    """Greedy non-maximum merging for axis-aligned bounding boxes.

    Instead of discarding overlapping boxes, merges them into the
    highest-scored box. Dispatches to the resolved backend.

    Args:
        predictions: Array of shape (N, 6) with columns
            [x1, y1, x2, y2, score, category_id].
        match_metric: Overlap metric, "IOU" or "IOS".
        match_threshold: Minimum overlap to merge a lower-scored box.

    Returns:
        Dict mapping each kept index to a list of indices merged into it.
    """
    return _dispatch("greedy_nmm")(predictions, match_metric, match_threshold)

def batched_greedy_nmm(
    predictions: np.ndarray,
    match_metric: str = "IOU",
    match_threshold: float = 0.5,
) -> dict[int, list[int]]:
    """Apply greedy non-maximum merging independently per category.

    Args:
        predictions: Array of shape (N, 6) with columns
            [x1, y1, x2, y2, score, category_id].
        match_metric: Overlap metric, "IOU" or "IOS".
        match_threshold: Minimum overlap to merge a lower-scored box.

    Returns:
        Dict mapping each kept index to a list of indices merged into it.
    """
    return _batched_apply(predictions, greedy_nmm, match_metric, match_threshold)  # type: ignore[return-value]

def nmm(
    predictions: np.ndarray,
    match_metric: str = "IOU",
    match_threshold: float = 0.5,
) -> dict[int, list[int]]:
    """Non-maximum merging (non-greedy, transitive) for axis-aligned bounding boxes.

    Unlike greedy NMM, this variant allows transitive merging: if box A
    merges with B and B merges with C, all three are merged together.
    Dispatches to the resolved backend.

    Args:
        predictions: Array of shape (N, 6) with columns
            [x1, y1, x2, y2, score, category_id].
        match_metric: Overlap metric, "IOU" or "IOS".
        match_threshold: Minimum overlap to merge a lower-scored box.

    Returns:
        Dict mapping each kept index to a list of indices merged into it.
    """
    return _dispatch("nmm")(predictions, match_metric, match_threshold)

def batched_nmm(
    predictions: np.ndarray,
    match_metric: str = "IOU",
    match_threshold: float = 0.5,
) -> dict[int, list[int]]:
    """Apply non-maximum merging (non-greedy, transitive) independently per category.

    Args:
        predictions: Array of shape (N, 6) with columns
            [x1, y1, x2, y2, score, category_id].
        match_metric: Overlap metric, "IOU" or "IOS".
        match_threshold: Minimum overlap to merge a lower-scored box.

    Returns:
        Dict mapping each kept index to a list of indices merged into it.
    """
    return _batched_apply(predictions, nmm, match_metric, match_threshold)  # type: ignore[return-value]