inference/core/nms.py

from typing import Optional

import numpy as np


def w_np_non_max_suppression(
    prediction,
    conf_thresh: float = 0.25,
    iou_thresh: float = 0.45,
    class_agnostic: bool = False,
    max_detections: int = 300,
    max_candidate_detections: int = 3000,
    timeout_seconds: Optional[int] = None,
    num_masks: int = 0,
    box_format: str = "xywh",
):
    """Applies non-maximum suppression to predictions.

    Args:
        prediction (np.ndarray): Array of predictions. Format for single prediction is
            [bbox x 4, max_class_confidence, (confidence) x num_of_classes, additional_element x num_masks]
        conf_thresh (float, optional): Confidence threshold. Defaults to 0.25.
        iou_thresh (float, optional): IOU threshold. Defaults to 0.45.
        class_agnostic (bool, optional): Whether to ignore class labels. Defaults to False.
        max_detections (int, optional): Maximum number of detections. Defaults to 300.
        max_candidate_detections (int, optional): Maximum number of candidate detections. Defaults to 3000.
        timeout_seconds (Optional[int], optional): Timeout in seconds. Defaults to None.
        num_masks (int, optional): Number of masks. Defaults to 0.
        box_format (str, optional): Format of bounding boxes. Either 'xywh' or 'xyxy'. Defaults to 'xywh'.

    Returns:
        list: List of filtered predictions after non-maximum suppression. Format of a single result is:
            [bbox x 4, max_class_confidence, max_class_confidence, id_of_class_with_max_confidence,
            additional_element x num_masks]
    """
    num_classes = prediction.shape[2] - 5 - num_masks

    np_box_corner = np.zeros(prediction.shape)
    if box_format == "xywh":
        np_box_corner[:, :, 0] = prediction[:, :, 0] - prediction[:, :, 2] / 2
        np_box_corner[:, :, 1] = prediction[:, :, 1] - prediction[:, :, 3] / 2
        np_box_corner[:, :, 2] = prediction[:, :, 0] + prediction[:, :, 2] / 2
        np_box_corner[:, :, 3] = prediction[:, :, 1] + prediction[:, :, 3] / 2
        prediction[:, :, :4] = np_box_corner[:, :, :4]
    elif box_format == "xyxy":
        pass
    else:
        raise ValueError(
            "box_format must be either 'xywh' or 'xyxy', got {}".format(box_format)
        )

    batch_predictions = []
    for np_image_i, np_image_pred in enumerate(prediction):
        filtered_predictions = []
        np_conf_mask = (np_image_pred[:, 4] >= conf_thresh).squeeze()

        np_image_pred = np_image_pred[np_conf_mask]
        if np_image_pred.shape[0] == 0:
            batch_predictions.append(filtered_predictions)
            continue
        np_class_conf = np.max(np_image_pred[:, 5 : num_classes + 5], 1)
        np_class_pred = np.argmax(np_image_pred[:, 5 : num_classes + 5], 1)
        np_class_conf = np.expand_dims(np_class_conf, axis=1)
        np_class_pred = np.expand_dims(np_class_pred, axis=1)
        np_mask_pred = np_image_pred[:, 5 + num_classes :]
        np_detections = np.append(
            np.append(
                np.append(np_image_pred[:, :5], np_class_conf, axis=1),
                np_class_pred,
                axis=1,
            ),
            np_mask_pred,
            axis=1,
        )

        np_unique_labels = np.unique(np_detections[:, 6])

        if class_agnostic:
            np_detections_class = sorted(
                np_detections, key=lambda row: row[4], reverse=True
            )
            filtered_predictions.extend(
                non_max_suppression_fast(np.array(np_detections_class), iou_thresh)
            )
        else:
            for c in np_unique_labels:
                np_detections_class = np_detections[np_detections[:, 6] == c]
                np_detections_class = sorted(
                    np_detections_class, key=lambda row: row[4], reverse=True
                )
                filtered_predictions.extend(
                    non_max_suppression_fast(np.array(np_detections_class), iou_thresh)
                )
        filtered_predictions = sorted(
            filtered_predictions, key=lambda row: row[4], reverse=True
        )
        batch_predictions.append(filtered_predictions[:max_detections])
    return batch_predictions


# Malisiewicz et al.
def non_max_suppression_fast(boxes, overlapThresh):
    """Applies non-maximum suppression to bounding boxes.

    Args:
        boxes (np.ndarray): Array of bounding boxes with confidence scores.
        overlapThresh (float): Overlap threshold for suppression.

    Returns:
        list: List of bounding boxes after non-maximum suppression.
    """
    # if there are no boxes, return an empty list
    if len(boxes) == 0:
        return []
    # if the bounding boxes integers, convert them to floats --
    # this is important since we'll be doing a bunch of divisions
    if boxes.dtype.kind == "i":
        boxes = boxes.astype("float")
    # initialize the list of picked indexes
    pick = []
    # grab the coordinates of the bounding boxes
    x1 = boxes[:, 0]
    y1 = boxes[:, 1]
    x2 = boxes[:, 2]
    y2 = boxes[:, 3]
    conf = boxes[:, 4]
    # compute the area of the bounding boxes and sort the bounding
    # boxes by the bottom-right y-coordinate of the bounding box
    area = (x2 - x1 + 1) * (y2 - y1 + 1)
    idxs = np.argsort(conf)
    # keep looping while some indexes still remain in the indexes
    # list
    while len(idxs) > 0:
        # grab the last index in the indexes list and add the
        # index value to the list of picked indexes
        last = len(idxs) - 1
        i = idxs[last]
        pick.append(i)
        # find the largest (x, y) coordinates for the start of
        # the bounding box and the smallest (x, y) coordinates
        # for the end of the bounding box
        xx1 = np.maximum(x1[i], x1[idxs[:last]])
        yy1 = np.maximum(y1[i], y1[idxs[:last]])
        xx2 = np.minimum(x2[i], x2[idxs[:last]])
        yy2 = np.minimum(y2[i], y2[idxs[:last]])
        # compute the width and height of the bounding box
        w = np.maximum(0, xx2 - xx1 + 1)
        h = np.maximum(0, yy2 - yy1 + 1)
        # compute the ratio of overlap
        overlap = (w * h) / area[idxs[:last]]
        # delete all indexes from the index list that have
        idxs = np.delete(
            idxs, np.concatenate(([last], np.where(overlap > overlapThresh)[0]))
        )
    # return only the bounding boxes that were picked using the
    # integer data type
    return boxes[pick].astype("float")