Merge pull request #4 from roboflow/feature/initial-line-counter-api

feature/initial-line-counter-api

supervision/tools/line_counter.py

@@ -0,0 +1,207 @@
+from typing import Dict
+
+import cv2
+import numpy as np
+
+from supervision.draw.color import Color
+from supervision.geometry.dataclasses import Point, Rect, Vector
+from supervision.tools.detections import Detections
+
+
+class LineCounter:
+    def __init__(self, start: Point, end: Point):
+        """
+        Initialize a LineCounter object.
+
+        :param start: Point : The starting point of the line.
+        :param end: Point : The ending point of the line.
+        """
+        self.vector = Vector(start=start, end=end)
+        self.tracker_state: Dict[str, bool] = {}
+        self.in_count: int = 0
+        self.out_count: int = 0
+
+    def update(self, detections: Detections):
+        """
+        Update the in_count and out_count for the detections that cross the line.
+
+        :param detections: Detections : The detections for which to update the counts.
+        """
+        for xyxy, confidence, class_id, tracker_id in detections:
+            # handle detections with no tracker_id
+            if tracker_id is None:
+                continue
+
+            # we check if all four anchors of bbox are on the same side of vector
+            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),
+            ]
+            triggers = [self.vector.is_in(point=anchor) for anchor in anchors]
+
+            # detection is partially in and partially out
+            if len(set(triggers)) == 2:
+                continue
+
+            tracker_state = triggers[0]
+            # handle new detection
+            if tracker_id not in self.tracker_state:
+                self.tracker_state[tracker_id] = tracker_state
+                continue
+
+            # handle detection on the same side of the line
+            if self.tracker_state.get(tracker_id) == tracker_state:
+                continue
+
+            self.tracker_state[tracker_id] = tracker_state
+            if tracker_state:
+                self.in_count += 1
+            else:
+                self.out_count += 1
+
+
+class LineCounterAnnotator:
+    def __init__(
+        self,
+        thickness: float = 2,
+        color: Color = Color.white(),
+        text_thickness: float = 2,
+        text_color: Color = Color.black(),
+        text_scale: float = 0.5,
+        text_offset: float = 1.5,
+        text_padding: int = 10,
+    ):
+        """
+        Initialize the LineCounterAnnotator object with default values.
+
+        :param thickness: float : The thickness of the line that will be drawn.
+        :param color: Color : The color of the line that will be drawn.
+        :param text_thickness: float : The thickness of the text that will be drawn.
+        :param text_color: Color : The color of the text that will be drawn.
+        :param text_scale: float : The scale of the text that will be drawn.
+        :param text_offset: float : The offset of the text that will be drawn.
+        :param text_padding: int : The padding of the text that will be drawn.
+        """
+        self.thickness: float = thickness
+        self.color: Color = color
+        self.text_thickness: float = text_thickness
+        self.text_color: Color = text_color
+        self.text_scale: float = text_scale
+        self.text_offset: float = text_offset
+        self.text_padding: int = text_padding
+
+    def annotate(self, frame: np.ndarray, line_counter: LineCounter) -> np.ndarray:
+        """
+        Draws the line on the frame using the line_counter provided.
+
+        :param frame: np.ndarray : The image on which the line will be drawn
+        :param line_counter: LineCounter : The line counter that will be used to draw the line
+        :return: np.ndarray : The image with the line drawn on it
+        """
+        cv2.line(
+            frame,
+            line_counter.vector.start.as_xy_int_tuple(),
+            line_counter.vector.end.as_xy_int_tuple(),
+            self.color.as_bgr(),
+            self.thickness,
+            lineType=cv2.LINE_AA,
+            shift=0,
+        )
+        cv2.circle(
+            frame,
+            line_counter.vector.start.as_xy_int_tuple(),
+            radius=5,
+            color=self.text_color.as_bgr(),
+            thickness=-1,
+            lineType=cv2.LINE_AA,
+        )
+        cv2.circle(
+            frame,
+            line_counter.vector.end.as_xy_int_tuple(),
+            radius=5,
+            color=self.text_color.as_bgr(),
+            thickness=-1,
+            lineType=cv2.LINE_AA,
+        )
+
+        in_text = f"in: {line_counter.in_count}"
+        out_text = f"out: {line_counter.out_count}"
+
+        (in_text_width, in_text_height), _ = cv2.getTextSize(
+            in_text, cv2.FONT_HERSHEY_SIMPLEX, self.text_scale, self.text_thickness
+        )
+        (out_text_width, out_text_height), _ = cv2.getTextSize(
+            out_text, cv2.FONT_HERSHEY_SIMPLEX, self.text_scale, self.text_thickness
+        )
+
+        in_text_x = int(
+            (line_counter.vector.end.x + line_counter.vector.start.x - in_text_width)
+            / 2
+        )
+        in_text_y = int(
+            (line_counter.vector.end.y + line_counter.vector.start.y + in_text_height)
+            / 2
+            - self.text_offset * in_text_height
+        )
+
+        out_text_x = int(
+            (line_counter.vector.end.x + line_counter.vector.start.x - out_text_width)
+            / 2
+        )
+        out_text_y = int(
+            (line_counter.vector.end.y + line_counter.vector.start.y + out_text_height)
+            / 2
+            + self.text_offset * out_text_height
+        )
+
+        in_text_background_rect = Rect(
+            x=in_text_x,
+            y=in_text_y - in_text_height,
+            width=in_text_width,
+            height=in_text_height,
+        ).pad(padding=self.text_padding)
+        out_text_background_rect = Rect(
+            x=out_text_x,
+            y=out_text_y - out_text_height,
+            width=out_text_width,
+            height=out_text_height,
+        ).pad(padding=self.text_padding)
+
+        cv2.rectangle(
+            frame,
+            in_text_background_rect.top_left.as_xy_int_tuple(),
+            in_text_background_rect.bottom_right.as_xy_int_tuple(),
+            self.color.as_bgr(),
+            -1,
+        )
+        cv2.rectangle(
+            frame,
+            out_text_background_rect.top_left.as_xy_int_tuple(),
+            out_text_background_rect.bottom_right.as_xy_int_tuple(),
+            self.color.as_bgr(),
+            -1,
+        )
+
+        cv2.putText(
+            frame,
+            in_text,
+            (in_text_x, in_text_y),
+            cv2.FONT_HERSHEY_SIMPLEX,
+            self.text_scale,
+            self.text_color.as_bgr(),
+            self.text_thickness,
+            cv2.LINE_AA,
+        )
+        cv2.putText(
+            frame,
+            out_text,
+            (out_text_x, out_text_y),
+            cv2.FONT_HERSHEY_SIMPLEX,
+            self.text_scale,
+            self.text_color.as_bgr(),
+            self.text_thickness,
+            cv2.LINE_AA,
+        )