diff --git a/README.md b/README.md
index a6688a957..8042fda01 100644
--- a/README.md
+++ b/README.md
@@ -7,6 +7,7 @@
>
+
@@ -53,6 +54,10 @@
+
+
+[](https://colab.research.google.com/github/roboflow-ai/notebooks/blob/main/notebooks/how-to-detect-and-count-objects-in-polygon-zone.ipynb)
+
## 👋 hello
diff --git a/supervision/__init__.py b/supervision/__init__.py
index ea2f155ea..1ae9bc686 100644
--- a/supervision/__init__.py
+++ b/supervision/__init__.py
@@ -1,8 +1,8 @@
-__version__ = "0.2.1"
+__version__ = "0.3.0"
from supervision.detection.core import BoxAnnotator, Detections
-from supervision.detection.polygon_zone import PolygonZone, PolygonZoneAnnotator
from supervision.detection.line_counter import LineZone, LineZoneAnnotator
+from supervision.detection.polygon_zone import PolygonZone, PolygonZoneAnnotator
from supervision.detection.utils import generate_2d_mask
from supervision.draw.color import Color, ColorPalette
from supervision.draw.utils import draw_filled_rectangle, draw_polygon, draw_text
diff --git a/supervision/detection/core.py b/supervision/detection/core.py
index cef6c577d..64f36907e 100644
--- a/supervision/detection/core.py
+++ b/supervision/detection/core.py
@@ -1,11 +1,12 @@
from __future__ import annotations
from dataclasses import dataclass
-from typing import List, Optional, Union
+from typing import Iterator, List, Optional, Tuple, Union
import cv2
import numpy as np
+from supervision.detection.utils import non_max_suppression
from supervision.draw.color import Color, ColorPalette
from supervision.geometry.core import Position
@@ -17,22 +18,26 @@ class Detections:
Attributes:
xyxy (np.ndarray): An array of shape `(n, 4)` containing the bounding boxes coordinates in format `[x1, y1, x2, y2]`
- confidence (np.ndarray): An array of shape `(n,)` containing the confidence scores of the detections.
+ confidence (Optional[np.ndarray]): An array of shape `(n,)` containing the confidence scores of the detections.
class_id (np.ndarray): An array of shape `(n,)` containing the class ids of the detections.
tracker_id (Optional[np.ndarray]): An array of shape `(n,)` containing the tracker ids of the detections.
"""
xyxy: np.ndarray
- confidence: np.ndarray
class_id: np.ndarray
+ confidence: Optional[np.ndarray] = None
tracker_id: Optional[np.ndarray] = None
def __post_init__(self):
n = len(self.xyxy)
validators = [
(isinstance(self.xyxy, np.ndarray) and self.xyxy.shape == (n, 4)),
- (isinstance(self.confidence, np.ndarray) and self.confidence.shape == (n,)),
(isinstance(self.class_id, np.ndarray) and self.class_id.shape == (n,)),
+ self.confidence is None
+ or (
+ isinstance(self.confidence, np.ndarray)
+ and self.confidence.shape == (n,)
+ ),
self.tracker_id is None
or (
isinstance(self.tracker_id, np.ndarray)
@@ -42,7 +47,7 @@ def __post_init__(self):
if not all(validators):
raise ValueError(
"xyxy must be 2d np.ndarray with (n, 4) shape, "
- "confidence must be 1d np.ndarray with (n,) shape, "
+ "confidence must be None or 1d np.ndarray with (n,) shape, "
"class_id must be 1d np.ndarray with (n,) shape, "
"tracker_id must be None or 1d np.ndarray with (n,) shape"
)
@@ -53,14 +58,16 @@ def __len__(self):
"""
return len(self.xyxy)
- def __iter__(self):
+ def __iter__(
+ self,
+ ) -> Iterator[Tuple[np.ndarray, Optional[float], int, Optional[Union[str, int]]]]:
"""
Iterates over the Detections object and yield a tuple of `(xyxy, confidence, class_id, tracker_id)` for each detection.
"""
for i in range(len(self.xyxy)):
yield (
self.xyxy[i],
- self.confidence[i],
+ self.confidence[i] if self.confidence is not None else None,
self.class_id[i],
self.tracker_id[i] if self.tracker_id is not None else None,
)
@@ -69,11 +76,17 @@ def __eq__(self, other: Detections):
return all(
[
np.array_equal(self.xyxy, other.xyxy),
- np.array_equal(self.confidence, other.confidence),
+ any(
+ [
+ self.confidence is None and other.confidence is None,
+ np.array_equal(self.confidence, other.confidence),
+ ]
+ ),
np.array_equal(self.class_id, other.class_id),
any(
[
self.tracker_id is None and other.tracker_id is None,
+ np.array_equal(self.tracker_id, other.tracker_id),
]
),
]
@@ -122,7 +135,7 @@ def from_yolov8(cls, yolov8_results):
>>> from supervision import Detections
>>> model = YOLO('yolov8s.pt')
- >>> results = model(frame)
+ >>> results = model(frame)[0]
>>> detections = Detections.from_yolov8(results)
```
"""
@@ -132,6 +145,36 @@ def from_yolov8(cls, yolov8_results):
class_id=yolov8_results.boxes.cls.cpu().numpy().astype(int),
)
+ @classmethod
+ def from_transformers(cls, transformers_results: dict):
+ return cls(
+ xyxy=transformers_results["boxes"].cpu().numpy(),
+ confidence=transformers_results["scores"].cpu().numpy(),
+ class_id=transformers_results["labels"].cpu().numpy().astype(int),
+ )
+
+ @classmethod
+ def from_detectron2(cls, detectron2_results):
+ return cls(
+ xyxy=detectron2_results["instances"].pred_boxes.tensor.cpu().numpy(),
+ confidence=detectron2_results["instances"].scores.cpu().numpy(),
+ class_id=detectron2_results["instances"]
+ .pred_classes.cpu()
+ .numpy()
+ .astype(int),
+ )
+
+ @classmethod
+ def from_coco_annotations(cls, coco_annotation: dict):
+ xyxy, class_id = [], []
+
+ for annotation in coco_annotation:
+ x_min, y_min, width, height = annotation["bbox"]
+ xyxy.append([x_min, y_min, x_min + width, y_min + height])
+ class_id.append(annotation["category_id"])
+
+ 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.
@@ -186,7 +229,9 @@ def get_anchor_coordinates(self, anchor: Position) -> np.ndarray:
raise ValueError(f"{anchor} is not supported.")
def __getitem__(self, index: np.ndarray) -> Detections:
- if isinstance(index, np.ndarray) and index.dtype == bool:
+ if isinstance(index, np.ndarray) and (
+ index.dtype == bool or index.dtype == int
+ ):
return Detections(
xyxy=self.xyxy[index],
confidence=self.confidence[index],
@@ -199,6 +244,17 @@ def __getitem__(self, index: np.ndarray) -> Detections:
f"Detections.__getitem__ not supported for index of type {type(index)}."
)
+ @property
+ def area(self) -> np.ndarray:
+ return (self.xyxy[:, 3] - self.xyxy[:, 1]) * (self.xyxy[:, 2] - self.xyxy[:, 0])
+
+ def with_nms(self, threshold: float = 0.5) -> Detections:
+ 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)
+ return self[indices]
+
class BoxAnnotator:
def __init__(
@@ -266,7 +322,7 @@ def annotate(
continue
text = (
- f"{confidence:0.2f}"
+ f"{class_id}"
if (labels is None or len(detections) != len(labels))
else labels[i]
)
diff --git a/supervision/detection/utils.py b/supervision/detection/utils.py
index a3a33490f..0bb2734d3 100644
--- a/supervision/detection/utils.py
+++ b/supervision/detection/utils.py
@@ -18,3 +18,40 @@ def generate_2d_mask(polygon: np.ndarray, resolution_wh: Tuple[int, int]) -> np.
mask = np.zeros((height, width), dtype=np.uint8)
cv2.fillPoly(mask, [polygon], color=1)
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]
+
+ 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)
+
+
+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