xiaoxue_admin
/
supervision
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

👋 initial commit

This commit is contained in:
SkalskiP 2023-03-07 22:35:09 +01:00
parent ac1658268c
commit 25c19b4422
3 changed files with 244 additions and 69 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

90
supervision/detection/core.py
View File

@ -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]

101
supervision/detection/utils.py
View File

@ -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()]

122
test/detection/test_utils.py Normal file
View File

@ -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)