Merge pull request #735 from revtheundead/add-center-point-crossed-condition-to-line-counter

SkalskiP · web-flow · commit c04edbcf85cb · 2024-01-18T17:52:13.000+01:00
Added an option to count detection upon the center point of the bounding box crossing the line counter
diff --git a/examples/speed_estimation/README.md b/examples/speed_estimation/README.md
@@ -11,10 +11,10 @@ supervision package for multiple tasks such as tracking, annotations, etc.
 
 https://github.com/roboflow/supervision/assets/26109316/d50118c1-2ae4-458d-915a-5d860fd36f71
 
-> [!IMPORTANT] 
+> [!IMPORTANT]
 > Adjust the [`SOURCE`](https://github.com/roboflow/supervision/blob/e32b05a636dab2ea1f39299e529c4b22b8baa8da/examples/speed_estimation/ultralytics_example.py#L10)
 > and [`TARGET`](https://github.com/roboflow/supervision/blob/e32b05a636dab2ea1f39299e529c4b22b8baa8da/examples/speed_estimation/ultralytics_example.py#L15)
-> configuration if you plan to run a speed estimation script on your video file. Those must be adjusted separately for each camera view. You can learn more 
+> configuration if you plan to run a speed estimation script on your video file. Those must be adjusted separately for each camera view. You can learn more
 > from our YouTube [tutorial](https://youtu.be/uWP6UjDeZvY).
 
 ## 💻 install
diff --git a/supervision/detection/line_counter.py b/supervision/detection/line_counter.py
@@ -1,11 +1,11 @@
-from typing import Dict, Optional, Tuple
+from typing import Dict, Iterable, Optional, Tuple
 
 import cv2
 import numpy as np
 
 from supervision.detection.core import Detections
 from supervision.draw.color import Color
-from supervision.geometry.core import Point, Rect, Vector
+from supervision.geometry.core import Point, Position, Rect, Vector
 
 
 class LineZone:
@@ -26,16 +26,71 @@ class LineZone:
             to outside.
     """
 
-    def __init__(self, start: Point, end: Point):
+    def __init__(
+        self,
+        start: Point,
+        end: Point,
+        triggering_anchors: Iterable[Position] = (
+            Position.TOP_LEFT,
+            Position.TOP_RIGHT,
+            Position.BOTTOM_LEFT,
+            Position.BOTTOM_RIGHT,
+        ),
+    ):
         """
         Args:
             start (Point): The starting point of the line.
             end (Point): The ending point of the line.
+            triggering_anchors (List[sv.Position]): A list of positions
+                specifying which anchors of the detections bounding box
+                to consider when deciding on whether the detection
+                has passed the line counter or not. By default, this
+                contains the four corners of the detection's bounding box
         """
         self.vector = Vector(start=start, end=end)
+        self.limits = self.calculate_region_of_interest_limits(vector=self.vector)
         self.tracker_state: Dict[str, bool] = {}
         self.in_count: int = 0
         self.out_count: int = 0
+        self.triggering_anchors = triggering_anchors
+
+    @staticmethod
+    def calculate_region_of_interest_limits(vector: Vector) -> Tuple[Vector, Vector]:
+        magnitude = vector.magnitude
+
+        if magnitude == 0:
+            raise ValueError("The magnitude of the vector cannot be zero.")
+
+        delta_x = vector.end.x - vector.start.x
+        delta_y = vector.end.y - vector.start.y
+
+        unit_vector_x = delta_x / magnitude
+        unit_vector_y = delta_y / magnitude
+
+        perpendicular_vector_x = -unit_vector_y
+        perpendicular_vector_y = unit_vector_x
+
+        start_region_limit = Vector(
+            start=vector.start,
+            end=Point(
+                x=vector.start.x + perpendicular_vector_x,
+                y=vector.start.y + perpendicular_vector_y,
+            ),
+        )
+        end_region_limit = Vector(
+            start=vector.end,
+            end=Point(
+                x=vector.end.x - perpendicular_vector_x,
+                y=vector.end.y - perpendicular_vector_y,
+            ),
+        )
+        return start_region_limit, end_region_limit
+
+    @staticmethod
+    def is_point_in_limits(point: Point, limits: Tuple[Vector, Vector]) -> bool:
+        cross_product_1 = limits[0].cross_product(point)
+        cross_product_2 = limits[1].cross_product(point)
+        return (cross_product_1 > 0) == (cross_product_2 > 0)
 
     def trigger(self, detections: Detections) -> Tuple[np.ndarray, np.ndarray]:
         """
@@ -54,20 +109,35 @@ def trigger(self, detections: Detections) -> Tuple[np.ndarray, np.ndarray]:
         crossed_in = np.full(len(detections), False)
         crossed_out = np.full(len(detections), False)
 
-        for i, (xyxy, tracker_id) in enumerate(
-            zip(detections.xyxy, detections.tracker_id)
-        ):
+        if len(detections) == 0:
+            return crossed_in, crossed_out
+
+        all_anchors = np.array(
+            [
+                detections.get_anchors_coordinates(anchor)
+                for anchor in self.triggering_anchors
+            ]
+        )
+
+        for i, tracker_id in enumerate(detections.tracker_id):
             if tracker_id is None:
                 continue
 
-            x1, y1, x2, y2 = xyxy
-            anchors = [
-                Point(x=x1, y=y1),
-                Point(x=x1, y=y2),
-                Point(x=x2, y=y1),
-                Point(x=x2, y=y2),
+            box_anchors = [Point(x=x, y=y) for x, y in all_anchors[:, i, :]]
+
+            in_limits = all(
+                [
+                    self.is_point_in_limits(point=anchor, limits=self.limits)
+                    for anchor in box_anchors
+                ]
+            )
+
+            if not in_limits:
+                continue
+
+            triggers = [
+                self.vector.cross_product(point=anchor) > 0 for anchor in box_anchors
             ]
-            triggers = [self.vector.is_in(point=anchor) for anchor in anchors]
 
             if len(set(triggers)) == 2:
                 continue
diff --git a/supervision/geometry/core.py b/supervision/geometry/core.py
@@ -2,6 +2,7 @@
 
 from dataclasses import dataclass
 from enum import Enum
+from math import sqrt
 from typing import Tuple
 
 
@@ -43,13 +44,38 @@ class Vector:
     start: Point
     end: Point
 
-    def is_in(self, point: Point) -> bool:
-        v1 = Vector(self.start, self.end)
-        v2 = Vector(self.start, point)
-        cross_product = (v1.end.x - v1.start.x) * (v2.end.y - v2.start.y) - (
-            v1.end.y - v1.start.y
-        ) * (v2.end.x - v2.start.x)
-        return cross_product < 0
+    @property
+    def magnitude(self) -> float:
+        """
+        Calculate the magnitude (length) of the vector.
+
+        Returns:
+            float: The magnitude of the vector.
+        """
+        dx = self.end.x - self.start.x
+        dy = self.end.y - self.start.y
+        return sqrt(dx**2 + dy**2)
+
+    def cross_product(self, point: Point) -> float:
+        """
+        Calculate the 2D cross product (also known as the vector product or outer
+        product) of the vector and a point, treated as vectors in 2D space.
+
+        Args:
+            point (Point): The point to be evaluated, treated as the endpoint of a
+                vector originating from the 'start' of the main vector.
+
+        Returns:
+            float: The scalar value of the cross product. It is positive if 'point'
+                lies to the left of the vector (when moving from 'start' to 'end'),
+                negative if it lies to the right, and 0 if it is collinear with the
+                vector.
+        """
+        dx_vector = self.end.x - self.start.x
+        dy_vector = self.end.y - self.start.y
+        dx_point = point.x - self.start.x
+        dy_point = point.y - self.start.y
+        return (dx_vector * dy_point) - (dy_vector * dx_point)
 
 
 @dataclass
diff --git a/test/detection/test_line_counter.py b/test/detection/test_line_counter.py
@@ -0,0 +1,72 @@
+from contextlib import ExitStack as DoesNotRaise
+from typing import Optional, Tuple
+
+import pytest
+
+from supervision import LineZone
+from supervision.geometry.core import Point, Vector
+
+
+@pytest.mark.parametrize(
+    "vector, expected_result, exception",
+    [
+        (
+            Vector(start=Point(x=0.0, y=0.0), end=Point(x=0.0, y=0.0)),
+            None,
+            pytest.raises(ValueError),
+        ),
+        (
+            Vector(start=Point(x=1.0, y=1.0), end=Point(x=1.0, y=1.0)),
+            None,
+            pytest.raises(ValueError),
+        ),
+        (
+            Vector(start=Point(x=0.0, y=0.0), end=Point(x=0.0, y=4.0)),
+            (
+                Vector(start=Point(x=0.0, y=0.0), end=Point(x=-1.0, y=0.0)),
+                Vector(start=Point(x=0.0, y=4.0), end=Point(x=1.0, y=4.0)),
+            ),
+            DoesNotRaise(),
+        ),
+        (
+            Vector(Point(0.0, 0.0), Point(4.0, 0.0)),
+            (
+                Vector(start=Point(x=0.0, y=0.0), end=Point(x=0.0, y=1.0)),
+                Vector(start=Point(x=4.0, y=0.0), end=Point(x=4.0, y=-1.0)),
+            ),
+            DoesNotRaise(),
+        ),
+        (
+            Vector(Point(0.0, 0.0), Point(3.0, 4.0)),
+            (
+                Vector(start=Point(x=0, y=0), end=Point(x=-0.8, y=0.6)),
+                Vector(start=Point(x=3, y=4), end=Point(x=3.8, y=3.4)),
+            ),
+            DoesNotRaise(),
+        ),
+        (
+            Vector(Point(0.0, 0.0), Point(4.0, 3.0)),
+            (
+                Vector(start=Point(x=0, y=0), end=Point(x=-0.6, y=0.8)),
+                Vector(start=Point(x=4, y=3), end=Point(x=4.6, y=2.2)),
+            ),
+            DoesNotRaise(),
+        ),
+        (
+            Vector(Point(0.0, 0.0), Point(3.0, -4.0)),
+            (
+                Vector(start=Point(x=0, y=0), end=Point(x=0.8, y=0.6)),
+                Vector(start=Point(x=3, y=-4), end=Point(x=2.2, y=-4.6)),
+            ),
+            DoesNotRaise(),
+        ),
+    ],
+)
+def test_calculate_region_of_interest_limits(
+    vector: Vector,
+    expected_result: Optional[Tuple[Vector, Vector]],
+    exception: Exception,
+) -> None:
+    with exception:
+        result = LineZone.calculate_region_of_interest_limits(vector=vector)
+        assert result == expected_result
diff --git a/test/geometry/test_core.py b/test/geometry/test_core.py
@@ -0,0 +1,59 @@
+import pytest
+
+from supervision.geometry.core import Point, Vector
+
+
+@pytest.mark.parametrize(
+    "vector, point, expected_result",
+    [
+        (Vector(start=Point(x=0, y=0), end=Point(x=5, y=5)), Point(x=-1, y=1), 10.0),
+        (Vector(start=Point(x=0, y=0), end=Point(x=5, y=5)), Point(x=6, y=6), 0.0),
+        (Vector(start=Point(x=0, y=0), end=Point(x=5, y=5)), Point(x=3, y=6), 15.0),
+        (Vector(start=Point(x=5, y=5), end=Point(x=0, y=0)), Point(x=-1, y=1), -10.0),
+        (Vector(start=Point(x=5, y=5), end=Point(x=0, y=0)), Point(x=6, y=6), 0.0),
+        (Vector(start=Point(x=5, y=5), end=Point(x=0, y=0)), Point(x=3, y=6), -15.0),
+        (Vector(start=Point(x=0, y=0), end=Point(x=1, y=0)), Point(x=0, y=0), 0.0),
+        (Vector(start=Point(x=0, y=0), end=Point(x=1, y=0)), Point(x=0, y=-1), -1.0),
+        (Vector(start=Point(x=0, y=0), end=Point(x=1, y=0)), Point(x=0, y=1), 1.0),
+        (Vector(start=Point(x=1, y=0), end=Point(x=0, y=0)), Point(x=0, y=0), 0.0),
+        (Vector(start=Point(x=1, y=0), end=Point(x=0, y=0)), Point(x=0, y=-1), 1.0),
+        (Vector(start=Point(x=1, y=0), end=Point(x=0, y=0)), Point(x=0, y=1), -1.0),
+        (Vector(start=Point(x=1, y=1), end=Point(x=1, y=3)), Point(x=0, y=0), 2.0),
+        (Vector(start=Point(x=1, y=1), end=Point(x=1, y=3)), Point(x=1, y=4), 0.0),
+        (Vector(start=Point(x=1, y=1), end=Point(x=1, y=3)), Point(x=2, y=4), -2.0),
+        (Vector(start=Point(x=1, y=3), end=Point(x=1, y=1)), Point(x=0, y=0), -2.0),
+        (Vector(start=Point(x=1, y=3), end=Point(x=1, y=1)), Point(x=1, y=4), 0.0),
+        (Vector(start=Point(x=1, y=3), end=Point(x=1, y=1)), Point(x=2, y=4), 2.0),
+    ],
+)
+def test_vector_cross_product(
+    vector: Vector, point: Point, expected_result: float
+) -> None:
+    result = vector.cross_product(point=point)
+    assert result == expected_result
+
+
+@pytest.mark.parametrize(
+    "vector, expected_result",
+    [
+        (Vector(start=Point(x=0, y=0), end=Point(x=0, y=0)), 0.0),
+        (Vector(start=Point(x=1, y=0), end=Point(x=0, y=0)), 1.0),
+        (Vector(start=Point(x=0, y=1), end=Point(x=0, y=0)), 1.0),
+        (Vector(start=Point(x=0, y=0), end=Point(x=1, y=0)), 1.0),
+        (Vector(start=Point(x=0, y=0), end=Point(x=0, y=1)), 1.0),
+        (Vector(start=Point(x=-1, y=0), end=Point(x=0, y=0)), 1.0),
+        (Vector(start=Point(x=0, y=-1), end=Point(x=0, y=0)), 1.0),
+        (Vector(start=Point(x=0, y=0), end=Point(x=-1, y=0)), 1.0),
+        (Vector(start=Point(x=0, y=0), end=Point(x=0, y=-1)), 1.0),
+        (Vector(start=Point(x=0, y=0), end=Point(x=3, y=4)), 5.0),
+        (Vector(start=Point(x=0, y=0), end=Point(x=-3, y=4)), 5.0),
+        (Vector(start=Point(x=0, y=0), end=Point(x=3, y=-4)), 5.0),
+        (Vector(start=Point(x=0, y=0), end=Point(x=-3, y=-4)), 5.0),
+        (Vector(start=Point(x=0, y=0), end=Point(x=4, y=3)), 5.0),
+        (Vector(start=Point(x=3, y=4), end=Point(x=0, y=0)), 5.0),
+        (Vector(start=Point(x=4, y=3), end=Point(x=0, y=0)), 5.0),
+    ],
+)
+def test_vector_magnitude(vector: Vector, expected_result: float) -> None:
+    result = vector.magnitude
+    assert result == expected_result
diff --git a/test/geometry/test_dataclasses.py b/test/geometry/test_dataclasses.py
