Combine¶

class GreedyNMMPostprocess(NMMPostprocess):
    """Postprocessor using Greedy Non-Maximum Merging (NMM).

    Similar to NMM but uses a greedy strategy: each kept prediction only
    merges boxes that directly overlap with it (no transitive merging).
    This is faster than full NMM and produces tighter merged boxes.
    """

    _agnostic_func = staticmethod(greedy_nmm)
    _batched_func = staticmethod(batched_greedy_nmm)

class LSNMSPostprocess(PostprocessPredictions):
    """Postprocessor using Locality-Sensitive NMS from the ``lsnms`` package.

    Uses a spatial index for fast neighbor lookup, making it efficient for
    large numbers of predictions. Only supports IoU metric (not IoS).
    Requires the ``lsnms`` package (``pip install lsnms>0.3.1``).

    Note:
        This postprocessor is experimental and not recommended for
        production use.
    """

    def __call__(self, object_predictions: list[ObjectPrediction]) -> list[ObjectPrediction]:
        """Apply Locality-Sensitive NMS to suppress overlapping predictions.

        Args:
            object_predictions: List of ObjectPrediction instances to suppress.

        Returns:
            List of suppressed ObjectPrediction instances.

        Raises:
            ModuleNotFoundError: If the lsnms package is not installed.
            NotImplementedError: If match_metric is not "IOU".
        """
        try:
            from lsnms import nms
        except ModuleNotFoundError:
            raise ModuleNotFoundError(
                'Please run "pip install lsnms>0.3.1" to install lsnms first for lsnms utilities.'
            )

        if self.match_metric == "IOS":
            raise NotImplementedError(f"match_metric={self.match_metric} is not supported for LSNMSPostprocess")

        logger.warning("LSNMSPostprocess is experimental and not recommended to use.")

        object_prediction_list = ObjectPredictionList(object_predictions)
        preds_np = object_prediction_list.tonumpy()

        boxes = preds_np[:, :4]
        scores = preds_np[:, 4]
        class_ids = preds_np[:, 5].astype("uint8")

        keep = nms(
            boxes, scores, iou_threshold=self.match_threshold, class_ids=None if self.class_agnostic else class_ids
        )

        selected = object_prediction_list[keep].tolist()
        if not isinstance(selected, list):
            selected = [selected]
        return selected

class NMMPostprocess(PostprocessPredictions):
    """Postprocessor using Non-Maximum Merging (NMM) with transitive merging.

    Instead of discarding overlapping detections, merges their bounding
    boxes, masks, and scores. Uses non-greedy transitive merging: if A
    overlaps B and B overlaps C, all three are merged even if A does not
    directly overlap C.
    """

    _agnostic_func = staticmethod(nmm)
    _batched_func = staticmethod(batched_nmm)

    def __call__(self, object_predictions: list[ObjectPrediction]) -> list[ObjectPrediction]:
        """Apply Non-Maximum Merging to merge overlapping predictions.

        Args:
            object_predictions: List of ObjectPrediction instances to merge.

        Returns:
            List of merged ObjectPrediction instances.
        """
        object_prediction_list = ObjectPredictionList(object_predictions)
        preds_np = object_prediction_list.tonumpy()
        func = self._agnostic_func if self.class_agnostic else self._batched_func
        keep_to_merge = func(preds_np, match_threshold=self.match_threshold, match_metric=self.match_metric)
        return _apply_merge(object_prediction_list, keep_to_merge, self.match_metric, self.match_threshold)

class NMSPostprocess(PostprocessPredictions):
    """Postprocessor using Non-Maximum Suppression (NMS).

    Keeps the highest-scored prediction among overlapping boxes and
    discards the rest. Does not merge bounding boxes or masks.
    """

    def __call__(self, object_predictions: list[ObjectPrediction]) -> list[ObjectPrediction]:
        """Apply Non-Maximum Suppression to suppress overlapping predictions.

        Args:
            object_predictions: List of ObjectPrediction instances to suppress.

        Returns:
            List of suppressed ObjectPrediction instances.
        """
        object_prediction_list = ObjectPredictionList(object_predictions)
        preds_np = object_prediction_list.tonumpy()
        func = nms if self.class_agnostic else batched_nms
        keep = func(preds_np, match_threshold=self.match_threshold, match_metric=self.match_metric)

        selected = object_prediction_list[keep].tolist()
        if not isinstance(selected, list):
            selected = [selected]
        return selected

class PostprocessPredictions(ABC):
    """Abstract base class for postprocessing object prediction lists.

    Subclasses implement a specific strategy (NMS, NMM, greedy NMM, etc.)
    to reduce overlapping detections produced by sliced inference.

    Args:
        match_threshold: Minimum overlap value (IoU or IoS) to consider
            two predictions as matching.
        match_metric: Overlap metric, "IOU" or "IOS".
        class_agnostic: If True, apply postprocessing across all
            categories. If False, apply per category independently.
    """

    def __init__(
        self,
        match_threshold: float = 0.5,
        match_metric: str = "IOU",
        class_agnostic: bool = True,
    ) -> None:
        """Initialize the postprocessor with configuration parameters.

        Args:
            match_threshold: Minimum overlap value (IoU or IoS) to consider
                two predictions as matching.
            match_metric: Overlap metric, "IOU" or "IOS".
            class_agnostic: If True, apply postprocessing across all
                categories. If False, apply per category independently.
        """
        self.match_threshold = match_threshold
        self.class_agnostic = class_agnostic
        self.match_metric = match_metric

    @abstractmethod
    def __call__(self, predictions: list[ObjectPrediction]) -> list[ObjectPrediction]:
        """Apply postprocessing to the list of predictions.

        Args:
            predictions: List of ObjectPrediction instances to postprocess.

        Returns:
            List of postprocessed ObjectPrediction instances.
        """
        pass