- ColorPalette.default() added

- Detections.from_yolov5() updated
- Detection.__getitem__ added
- get_polygon_center added

supervision/__init__.py

@@ -1 +1,3 @@
 __version__ = "0.2.0.dev0"
+
+from supervision.detection.core import Detections

supervision/detection/core.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 from dataclasses import dataclass
 from typing import List, Optional, Union
 
@@ -76,7 +78,7 @@ class Detections:
         Example:
             ```python
             >>> import torch
-            >>> from supervision.detection.core import Detections
+            >>> from supervision import Detections
 
             >>> model = torch.hub.load('ultralytics/yolov5', 'yolov5s')
             >>> results = model(frame)
@@ -89,7 +91,7 @@ class Detections:
         class_id = yolov5_detections_predictions[:, 5].astype(int)
         return cls(xyxy, confidence, class_id)
 
-    def filter(self, mask: np.ndarray, inplace: bool = False) -> Optional[np.ndarray]:
+    def filter(self, mask: np.ndarray, inplace: bool = False) -> Optional[Detections]:
         """
         Filter the detections by applying a mask.
 
@@ -142,11 +144,21 @@ 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:
+            return Detections(
+                xyxy=self.xyxy[index],
+                confidence=self.confidence[index],
+                class_id=self.class_id[index],
+                tracker_id=self.tracker_id[index] if self.tracker_id is not None else None,
+            )
+        raise TypeError(f"Detections.__getitem__ not supported for index of type {type(index)}.")
+
 
 class BoxAnnotator:
     def __init__(
         self,
-        color: Union[Color, ColorPalette],
+        color: Union[Color, ColorPalette] = ColorPalette.default(),
         thickness: int = 2,
         text_color: Color = Color.black(),
         text_scale: float = 0.5,

supervision/draw/color.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from dataclasses import dataclass, field
+from dataclasses import dataclass
 from typing import List, Tuple
 
 DEFAULT_COLOR_PALETTE = [
@@ -94,11 +94,11 @@ class Color:
 
 @dataclass
 class ColorPalette:
-    colors: List[Color] = field(
-        default_factory=lambda: [
-            Color.from_hex(color_hex) for color_hex in DEFAULT_COLOR_PALETTE
-        ]
-    )
+    colors: List[Color]
+
+    @classmethod
+    def default(cls) -> ColorPalette:
+        return ColorPalette.from_hex(color_hex_list=DEFAULT_COLOR_PALETTE)
 
     @classmethod
     def from_hex(cls, color_hex_list: List[str]):

supervision/geometry/utils.py

@@ -0,0 +1,28 @@
+import numpy as np
+
+from supervision.geometry.core import Point
+
+
+def get_polygon_center(polygon: np.ndarray) -> Point:
+    """
+    Calculate the center of a polygon.
+
+    This function takes in a polygon as a 2-dimensional numpy ndarray and returns the center of the polygon as a Point object. The center is calculated as the mean of the polygon's vertices along each axis, and is rounded down to the nearest integer.
+
+    Parameters:
+        polygon (np.ndarray): A 2-dimensional numpy ndarray representing the vertices of the polygon.
+
+    Returns:
+        Point: The center of the polygon, represented as a Point object with x and y attributes.
+
+    Examples:
+        ```python
+        >>> from supervision.geometry.utils import get_polygon_center
+
+        >>> vertices = np.array([[0, 0], [0, 1], [1, 1], [1, 0]])
+        >>> get_center(vertices)
+        Point(x=0.5, y=0.5)
+        ```
+    """
+    center = np.mean(polygon, axis=0).astype(int)
+    return Point(x=center[0], y=center[1])