From 25c19b4422897a69cbf60addb4159ea03a9097e3 Mon Sep 17 00:00:00 2001 From: SkalskiP Date: Tue, 7 Mar 2023 22:35:09 +0100 Subject: [PATCH] =?UTF-8?q?=F0=9F=91=8B=20initial=20commit?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- supervision/detection/core.py | 90 ++++++++++++++---------- supervision/detection/utils.py | 101 ++++++++++++++++++--------- test/detection/test_utils.py | 122 +++++++++++++++++++++++++++++++++ 3 files changed, 244 insertions(+), 69 deletions(-) create mode 100644 test/detection/test_utils.py diff --git a/supervision/detection/core.py b/supervision/detection/core.py index 64f36907..545d6945 100644 --- a/supervision/detection/core.py +++ b/supervision/detection/core.py @@ -93,14 +93,15 @@ class Detections: ) @classmethod - def from_yolov5(cls, yolov5_detections): + def from_yolov5(cls, yolov5_results): """ Creates a Detections instance from a YOLOv5 output Detections - Attributes: - yolov5_detections (yolov5.models.common.Detections): The output Detections instance from YOLOv5 + Args: + yolov5_results (yolov5.models.common.Detections): The output Detections instance from YOLOv5 Returns: + Detections: A new Detections object. Example: ```python @@ -112,7 +113,7 @@ class Detections: >>> detections = Detections.from_yolov5(results) ``` """ - yolov5_detections_predictions = yolov5_detections.pred[0].cpu().cpu().numpy() + yolov5_detections_predictions = yolov5_results.pred[0].cpu().cpu().numpy() return cls( xyxy=yolov5_detections_predictions[:, :4], confidence=yolov5_detections_predictions[:, 4], @@ -124,10 +125,11 @@ class Detections: """ Creates a Detections instance from a YOLOv8 output Results - Attributes: + Args: yolov8_results (ultralytics.yolo.engine.results.Results): The output Results instance from YOLOv8 Returns: + Detections: A new Detections object. Example: ```python @@ -147,6 +149,12 @@ class Detections: @classmethod def from_transformers(cls, transformers_results: dict): + """ + Creates a Detections instance from Object Detection Transformer output Results + + Returns: + Detections: A new Detections object. + """ return cls( xyxy=transformers_results["boxes"].cpu().numpy(), confidence=transformers_results["scores"].cpu().numpy(), @@ -175,40 +183,11 @@ class Detections: return cls(xyxy=np.array(xyxy), class_id=np.array(class_id)) - def filter(self, mask: np.ndarray, inplace: bool = False) -> Optional[Detections]: - """ - Filter the detections by applying a mask. - - Attributes: - mask (np.ndarray): A mask of shape `(n,)` containing a boolean value for each detection indicating if it should be included in the filtered detections - inplace (bool): If True, the original data will be modified and self will be returned. - - Returns: - Optional[np.ndarray]: A new instance of Detections with the filtered detections, if inplace is set to `False`. `None` otherwise. - """ - if inplace: - self.xyxy = self.xyxy[mask] - self.confidence = self.confidence[mask] - self.class_id = self.class_id[mask] - self.tracker_id = ( - self.tracker_id[mask] if self.tracker_id is not None else None - ) - return self - else: - return Detections( - xyxy=self.xyxy[mask], - confidence=self.confidence[mask], - class_id=self.class_id[mask], - tracker_id=self.tracker_id[mask] - if self.tracker_id is not None - else None, - ) - def get_anchor_coordinates(self, anchor: Position) -> np.ndarray: """ Returns the bounding box coordinates for a specific anchor. - Properties: + Args: anchor (Position): Position of bounding box anchor for which to return the coordinates. Returns: @@ -246,13 +225,50 @@ class Detections: @property def area(self) -> np.ndarray: + """ + Calculate the area of each bounding box in the set of object detections. + + Returns: + np.ndarray: An array of floats containing the area of each bounding box in the format of (area_1, area_2, ..., area_n), where n is the number of detections. + """ return (self.xyxy[:, 3] - self.xyxy[:, 1]) * (self.xyxy[:, 2] - self.xyxy[:, 0]) - def with_nms(self, threshold: float = 0.5) -> Detections: + def with_nms( + self, threshold: float = 0.5, class_agnostic: bool = False + ) -> Detections: + """ + Perform non-maximum suppression on the current set of object detections. + + Args: + threshold (float, optional): The intersection-over-union threshold to use for non-maximum suppression. Defaults to 0.5. + class_agnostic (bool, optional): Whether to perform class-agnostic non-maximum suppression. If True, the class_id of each detection will be ignored. Defaults to False. + + Returns: + Detections: A new Detections object containing the subset of detections after non-maximum suppression. + + Raises: + AssertionError: If `confidence` is None and class_agnostic is False. If `class_id` is None and class_agnostic is False. + """ assert ( self.confidence is not None ), f"Detections confidence must be given for NMS to be executed." - indices = non_max_suppression(self.xyxy, self.confidence, threshold=threshold) + + if class_agnostic: + predictions = np.hstack((self.xyxy, self.confidence.reshape(-1, 1))) + indices = non_max_suppression( + predictions=predictions, iou_threshold=threshold + ) + return self[indices] + + assert self.class_id is not None, ( + f"Detections class_id must be given for NMS to be executed. If you intended to perform class agnostic " + f"NMS set class_agnostic=True." + ) + + predictions = np.hstack( + (self.xyxy, self.confidence.reshape(-1, 1), self.class_id.reshape(-1, 1)) + ) + indices = non_max_suppression(predictions=predictions, iou_threshold=threshold) return self[indices] diff --git a/supervision/detection/utils.py b/supervision/detection/utils.py index 0bb2734d..e4c0f385 100644 --- a/supervision/detection/utils.py +++ b/supervision/detection/utils.py @@ -20,38 +20,75 @@ def generate_2d_mask(polygon: np.ndarray, resolution_wh: Tuple[int, int]) -> np. return mask -def non_max_suppression(boxes: np.ndarray, scores: np.ndarray, threshold: float): - assert boxes.shape[0] == scores.shape[0] - ys1 = boxes[:, 0] - xs1 = boxes[:, 1] - ys2 = boxes[:, 2] - xs2 = boxes[:, 3] +def box_iou_batch(boxes_true: np.ndarray, boxes_detection: np.ndarray) -> np.ndarray: + """ + Compute Intersection over Union 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. - areas = (ys2 - ys1) * (xs2 - xs1) - scores_indexes = scores.argsort().tolist() - boxes_keep_index = [] - while len(scores_indexes): - index = scores_indexes.pop() - boxes_keep_index.append(index) - if not len(scores_indexes): - break - iou = compute_iou( - boxes[index], boxes[scores_indexes], areas[index], areas[scores_indexes] - ) - filtered_indexes = set((iou > threshold).nonzero()[0]) - scores_indexes = [ - v for (i, v) in enumerate(scores_indexes) if i not in filtered_indexes - ] - return np.array(boxes_keep_index) + Properties: + 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) -def compute_iou(box, boxes, box_area, boxes_area): - assert boxes.shape[0] == boxes_area.shape[0] - ys1 = np.maximum(box[0], boxes[:, 0]) - xs1 = np.maximum(box[1], boxes[:, 1]) - ys2 = np.minimum(box[2], boxes[:, 2]) - xs2 = np.minimum(box[3], boxes[:, 3]) - intersections = np.maximum(ys2 - ys1, 0) * np.maximum(xs2 - xs1, 0) - unions = box_area + boxes_area - intersections - iou = intersections / unions - return iou +def non_max_suppression( + predictions: np.ndarray, iou_threshold: float = 0.5 +) -> np.ndarray: + """ + Perform non-maximum suppression 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. Defaults to 0.5. + + 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()] diff --git a/test/detection/test_utils.py b/test/detection/test_utils.py new file mode 100644 index 00000000..9ddf3cc5 --- /dev/null +++ b/test/detection/test_utils.py @@ -0,0 +1,122 @@ +from contextlib import ExitStack as DoesNotRaise +from typing import Optional + +import pytest + +import numpy as np + +from supervision.detection.utils import non_max_suppression + + +@pytest.mark.parametrize( + "predictions, iou_threshold, expected_result, exception", + [ + ( + np.array([ + [10.0, 10.0, 40.0, 40.0, 0.8] + ]), + 0.5, + np.array([ + True + ]), + DoesNotRaise() + ), # single box with no category + ( + np.array([ + [10.0, 10.0, 40.0, 40.0, 0.8, 0] + ]), + 0.5, + np.array([ + True + ]), + DoesNotRaise() + ), # single box with category + ( + np.array([ + [10.0, 10.0, 40.0, 40.0, 0.8], + [15.0, 15.0, 40.0, 40.0, 0.9], + ]), + 0.5, + np.array([ + False, + True + ]), + DoesNotRaise() + ), # two boxes with no category + ( + np.array([ + [10.0, 10.0, 40.0, 40.0, 0.8, 0], + [15.0, 15.0, 40.0, 40.0, 0.9, 1], + ]), + 0.5, + np.array([ + True, + True + ]), + DoesNotRaise() + ), # two boxes with different category +( + np.array([ + [10.0, 10.0, 40.0, 40.0, 0.8, 0], + [15.0, 15.0, 40.0, 40.0, 0.9, 0], + ]), + 0.5, + np.array([ + True, + True + ]), + DoesNotRaise() + ), # two boxes with same category + ( + np.array([ + [0.0, 0.0, 30.0, 40.0, 0.8], + [5.0, 5.0, 35.0, 45.0, 0.9], + [10.0, 10.0, 40.0, 50.0, 0.85], + ]), + 0.5, + np.array([ + False, + True, + False + ]), + DoesNotRaise() + ), # three boxes with no category + ( + np.array([ + [0.0, 0.0, 30.0, 40.0, 0.8, 0], + [5.0, 5.0, 35.0, 45.0, 0.9, 1], + [10.0, 10.0, 40.0, 50.0, 0.85, 2], + ]), + 0.5, + np.array([ + True, + True, + True + ]), + DoesNotRaise() + ), # three boxes with same category + ( + np.array([ + [0.0, 0.0, 30.0, 40.0, 0.8, 0], + [5.0, 5.0, 35.0, 45.0, 0.9, 0], + [10.0, 10.0, 40.0, 50.0, 0.85, 1], + ]), + 0.5, + np.array([ + False, + True, + True + ]), + DoesNotRaise() + ), # three boxes with different category + ] +) +def test_non_max_suppression( + predictions: np.ndarray, + iou_threshold: float, + expected_result: Optional[np.ndarray], + exception: Exception +) -> None: + with exception: + result = non_max_suppression(predictions=predictions, iou_threshold=iou_threshold) + np.array_equal(result, expected_result)