Merge pull request #32 from roboflow/feature/detections_api_refinement

feature/detections_api_refinement

README.md

@@ -7,6 +7,7 @@
       >
     </a>
   </p>
+
   <br>
 
   <div align="center">
@@ -53,6 +54,10 @@
       </a>
   </div>
 
+  <br>
+
+[![Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/roboflow-ai/notebooks/blob/main/notebooks/how-to-detect-and-count-objects-in-polygon-zone.ipynb)
+
 </div>
 
 ## 👋 hello
@@ -97,9 +102,15 @@ Curious how Supervision can help you solve problems on your project? Visit our [
 Learn how to use YOLOv8, ByteTrack and **Supervision** to detect, track and count objects. 🔥
 [Subscribe](https://www.youtube.com/@Roboflow), and stay up to date with our latest YouTube videos!
 
+<p align="center">
+    <a href="https://youtu.be/l_kf9CfZ_8M">
+        <img src="https://user-images.githubusercontent.com/26109316/217950212-311de186-1862-4b4c-a86e-89cafd68b233.jpg" alt="count-objects-in-zone">
+    </a>
+</p>
+
 <p align="center">
     <a href="https://youtu.be/OS5qI9YBkfk">
-        <img src="https://user-images.githubusercontent.com/26109316/213702005-ddd568f0-b902-46c2-9af9-e6ff33db23bf.jpg" alt="latest-roboflow-tutorial">
+        <img src="https://user-images.githubusercontent.com/26109316/213702005-ddd568f0-b902-46c2-9af9-e6ff33db23bf.jpg" alt="count-objects-crossing-line">
     </a>
 </p>
 

requirements.txt

@@ -1,30 +0,0 @@
-certifi==2022.12.7
-charset-normalizer==3.0.1
-click==8.1.3
-colorama==0.4.6
-ghp-import==2.1.0
-griffe==0.25.4
-idna==3.4
-importlib-metadata==6.0.0
-Jinja2==3.1.2
-Markdown==3.3.7
-MarkupSafe==2.1.2
-mergedeep==1.3.4
-mkdocs==1.4.2
-mkdocs-autorefs==0.4.1
-mkdocs-material==9.0.9
-mkdocs-material-extensions==1.1.1
-mkdocstrings==0.20.0
-mkdocstrings-python==0.8.3
-packaging==23.0
-Pygments==2.14.0
-pymdown-extensions==9.9.2
-python-dateutil==2.8.2
-PyYAML==6.0
-pyyaml_env_tag==0.1
-regex==2022.10.31
-requests==2.28.2
-six==1.16.0
-urllib3==1.26.14
-watchdog==2.2.1
-zipp==3.12.0

setup.py

@@ -24,7 +24,7 @@ setuptools.setup(
     long_description_content_type='text/markdown',
     url='https://github.com/roboflow/supervision',
     install_requires=[
-        'numpy',
+        'numpy>=1.20.0',
         'opencv-python',
         'matplotlib'
     ],

supervision/__init__.py

@@ -1,6 +1,7 @@
-__version__ = "0.2.0"
+__version__ = "0.3.0"
 
 from supervision.detection.core import BoxAnnotator, Detections
+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

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 @@ class Detections:
         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,18 +58,40 @@ class Detections:
         """
         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,
             )
 
+    def __eq__(self, other: Detections):
+        return all(
+            [
+                np.array_equal(self.xyxy, other.xyxy),
+                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),
+                    ]
+                ),
+            ]
+        )
+
     @classmethod
     def from_yolov5(cls, yolov5_detections):
         """
@@ -108,7 +135,7 @@ class Detections:
             >>> from supervision import Detections
 
             >>> model = YOLO('yolov8s.pt')
-            >>> results = model(frame)
+            >>> results = model(frame)[0]
             >>> detections = Detections.from_yolov8(results)
             ```
         """
@@ -118,6 +145,36 @@ class Detections:
             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.
@@ -172,7 +229,9 @@ class Detections:
         raise ValueError(f"{anchor} is not supported.")
 
     def __getitem__(self, index: np.ndarray) -> Detections:
-        if isinstance(index, np.ndarray) and index.dtype == np.bool:
+        if isinstance(index, np.ndarray) and (
+            index.dtype == bool or index.dtype == int
+        ):
             return Detections(
                 xyxy=self.xyxy[index],
                 confidence=self.confidence[index],
@@ -185,6 +244,17 @@ class 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__(
@@ -252,7 +322,7 @@ class BoxAnnotator:
                 continue
 
             text = (
-                f"{confidence:0.2f}"
+                f"{class_id}"
                 if (labels is None or len(detections) != len(labels))
                 else labels[i]
             )

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

test/detection/test_core.py

@@ -0,0 +1,105 @@
+from contextlib import ExitStack as DoesNotRaise
+
+import pytest
+
+from supervision import Detections
+
+from typing import Optional, Union
+
+import numpy as np
+
+
+PREDICTIONS = np.array([
+        [       2254,         906,        2447,        1353,     0.90538,           0],
+        [       2049,        1133,        2226,        1371,     0.59002,          56],
+        [        727,        1224,         838,        1601,     0.51119,          39],
+        [        808,        1214,         910,        1564,     0.45287,          39],
+        [          6,          52,        1131,        2133,     0.45057,          72],
+        [        299,        1225,         512,        1663,     0.45029,          39],
+        [        529,         874,         645,         945,     0.31101,          39],
+        [          8,          47,        1935,        2135,     0.28192,          72],
+        [       2265,         813,        2328,         901,      0.2714,          62]
+    ], dtype=np.float32)
+
+DETECTIONS = Detections(
+    xyxy=PREDICTIONS[:, :4],
+    confidence=PREDICTIONS[:, 4],
+    class_id=PREDICTIONS[:, 5].astype(int)
+)
+
+
+PREDICTIONS = np.array([
+        [       2254,         906,        2447,        1353,     0.90538,           0],
+        [       2049,        1133,        2226,        1371,     0.59002,          56],
+        [        727,        1224,         838,        1601,     0.51119,          39],
+    ], dtype=np.float32)
+
+
+@pytest.mark.parametrize(
+    'detections, index, expected_result, exception',
+    [
+        (
+            DETECTIONS,
+            DETECTIONS.class_id == 0,
+            Detections(
+                xyxy=np.array([
+                    [       2254,         906,        2447,        1353]
+                ], dtype=np.float32),
+                confidence=np.array([
+                    0.90538
+                ], dtype=np.float32),
+                class_id=np.array([
+                    0
+                ], dtype=int)
+            ),
+            DoesNotRaise()
+        ),  # take only detections with class_id = 0
+        (
+            DETECTIONS,
+            DETECTIONS.confidence > 0.5,
+            Detections(
+                xyxy=np.array([
+                    [       2254,         906,        2447,        1353],
+                    [       2049,        1133,        2226,        1371],
+                    [        727,        1224,         838,        1601]
+                ], dtype=np.float32),
+                confidence=np.array([
+                    0.90538,
+                    0.59002,
+                    0.51119
+                ], dtype=np.float32),
+                class_id=np.array([
+                    0,
+                    56,
+                    39
+                ], dtype=int)
+            ),
+            DoesNotRaise()
+        ),  # take only detections with confidence > 0.5
+        (
+            DETECTIONS,
+            np.array([True, True, True, True, True, True, True, True, True], dtype=bool),
+            DETECTIONS,
+            DoesNotRaise()
+        ),  # take all detections
+        (
+            DETECTIONS,
+            np.array([False, False, False, False, False, False, False, False, False], dtype=bool),
+            Detections(
+                xyxy=np.empty((0, 4), dtype=np.float32),
+                confidence=np.array([], dtype=np.float32),
+                class_id=np.array([], dtype=int)
+            ),
+            DoesNotRaise()
+        ),  # take no detections
+    ]
+)
+def test_getitem(
+        detections: Detections,
+        index: Union[int, slice, np.ndarray],
+        expected_result: Optional[Detections],
+        exception: Exception
+) -> None:
+    with exception:
+        result = detections[index]
+        assert result == expected_result