Utils

generate_2d_mask

Generate a 2D mask from a polygon.

Parameters:

Name	Type	Description	Default
polygon	ndarray	The polygon for which the mask should be generated, given as a list of vertices.	required
resolution_wh	Tuple[int, int]	The width and height of the desired resolution.	required

Returns:

Type	Description
ndarray	np.ndarray: The generated 2D mask, where the polygon is marked with 1's and the rest is filled with 0's.

Source code in supervision/detection/utils.py

def generate_2d_mask(polygon: np.ndarray, resolution_wh: Tuple[int, int]) -> np.ndarray:
    """Generate a 2D mask from a polygon.

    Args:
        polygon (np.ndarray): The polygon for which the mask should be generated, given as a list of vertices.
        resolution_wh (Tuple[int, int]): The width and height of the desired resolution.

    Returns:
        np.ndarray: The generated 2D mask, where the polygon is marked with `1`'s and the rest is filled with `0`'s.
    """
    width, height = resolution_wh
    mask = np.zeros((height, width), dtype=np.uint8)
    cv2.fillPoly(mask, [polygon], color=1)
    return mask

box_iou_batch

Compute Intersection over Union (IoU) of two sets of bounding boxes - boxes_true and boxes_detection. Both sets of boxes are expected to be in (x_min, y_min, x_max, y_max) format.

Parameters:

Name	Type	Description	Default
boxes_true	ndarray	2D np.ndarray representing ground-truth boxes. shape = (N, 4) where N is number of true objects.	required
boxes_detection	ndarray	2D np.ndarray representing detection boxes. shape = (M, 4) where M is number of detected objects.	required

Returns:

Type	Description
ndarray	np.ndarray: Pairwise IoU of boxes from boxes_true and boxes_detection. shape = (N, M) where N is number of true objects and M is number of detected objects.

Source code in supervision/detection/utils.py

def box_iou_batch(boxes_true: np.ndarray, boxes_detection: np.ndarray) -> np.ndarray:
    """
    Compute Intersection over Union (IoU) of two sets of bounding boxes - `boxes_true` and `boxes_detection`. Both sets
    of boxes are expected to be in `(x_min, y_min, x_max, y_max)` format.

    Args:
        boxes_true (np.ndarray): 2D `np.ndarray` representing ground-truth boxes. `shape = (N, 4)` where `N` is number of true objects.
        boxes_detection (np.ndarray): 2D `np.ndarray` representing detection boxes. `shape = (M, 4)` where `M` is number of detected objects.

    Returns:
        np.ndarray: Pairwise IoU of boxes from `boxes_true` and `boxes_detection`. `shape = (N, M)` where `N` is number of true objects and `M` is number of detected objects.
    """

    def box_area(box):
        return (box[2] - box[0]) * (box[3] - box[1])

    area_true = box_area(boxes_true.T)
    area_detection = box_area(boxes_detection.T)

    top_left = np.maximum(boxes_true[:, None, :2], boxes_detection[:, :2])
    bottom_right = np.minimum(boxes_true[:, None, 2:], boxes_detection[:, 2:])

    area_inter = np.prod(np.clip(bottom_right - top_left, a_min=0, a_max=None), 2)
    return area_inter / (area_true[:, None] + area_detection - area_inter)

non_max_suppression

Perform Non-Maximum Suppression (NMS) on object detection predictions.

Parameters:

Name	Type	Description	Default
predictions	ndarray	An array of object detection predictions in the format of (x_min, y_min, x_max, y_max, score) or (x_min, y_min, x_max, y_max, score, class).	required
iou_threshold	float	The intersection-over-union threshold to use for non-maximum suppression.	0.5

Returns:

Type	Description
ndarray	np.ndarray: A boolean array indicating which predictions to keep after non-maximum suppression.

Raises:

Type	Description
AssertionError	If iou_threshold is not within the closed range from 0 to 1.

Source code in supervision/detection/utils.py

def non_max_suppression(
    predictions: np.ndarray, iou_threshold: float = 0.5
) -> np.ndarray:
    """
    Perform Non-Maximum Suppression (NMS) on object detection predictions.

    Args:
        predictions (np.ndarray): An array of object detection predictions in the format of `(x_min, y_min, x_max, y_max, score)` or `(x_min, y_min, x_max, y_max, score, class)`.
        iou_threshold (float, optional): The intersection-over-union threshold to use for non-maximum suppression.

    Returns:
        np.ndarray: A boolean array indicating which predictions to keep after non-maximum suppression.

    Raises:
        AssertionError: If `iou_threshold` is not within the closed range from `0` to `1`.
    """
    assert 0 <= iou_threshold <= 1, (
        f"Value of `iou_threshold` must be in the closed range from 0 to 1, "
        f"{iou_threshold} given."
    )
    rows, columns = predictions.shape

    # add column #5 - category filled with zeros for agnostic nms
    if columns == 5:
        predictions = np.c_[predictions, np.zeros(rows)]

    # sort predictions column #4 - score
    sort_index = np.flip(predictions[:, 4].argsort())
    predictions = predictions[sort_index]

    boxes = predictions[:, :4]
    categories = predictions[:, 5]
    ious = box_iou_batch(boxes, boxes)
    ious = ious - np.eye(rows)

    keep = np.ones(rows, dtype=bool)

    for index, (iou, category) in enumerate(zip(ious, categories)):
        if not keep[index]:
            continue

        # drop detections with iou > iou_threshold and same category as current detections
        condition = (iou > iou_threshold) & (categories == category)
        keep = keep & ~condition

    return keep[sort_index.argsort()]

mask_to_xyxy

Converts a 3D np.array of 2D bool masks into a 2D np.array of bounding boxes.

Parameters:

Name	Type	Description	Default
masks	ndarray	A 3D np.array of shape (N, W, H) containing 2D bool masks	required

Returns:

Type	Description
ndarray	np.ndarray: A 2D np.array of shape (N, 4) containing the bounding boxes (x_min, y_min, x_max, y_max) for each mask

Source code in supervision/detection/utils.py

def mask_to_xyxy(masks: np.ndarray) -> np.ndarray:
    """
    Converts a 3D `np.array` of 2D bool masks into a 2D `np.array` of bounding boxes.

    Parameters:
        masks (np.ndarray): A 3D `np.array` of shape `(N, W, H)` containing 2D bool masks

    Returns:
        np.ndarray: A 2D `np.array` of shape `(N, 4)` containing the bounding boxes `(x_min, y_min, x_max, y_max)` for each mask
    """
    n = masks.shape[0]
    bboxes = np.zeros((n, 4), dtype=int)

    for i, mask in enumerate(masks):
        rows, cols = np.where(mask)

        if len(rows) > 0 and len(cols) > 0:
            x_min, x_max = np.min(cols), np.max(cols)
            y_min, y_max = np.min(rows), np.max(rows)
            bboxes[i, :] = [x_min, y_min, x_max, y_max]

    return bboxes