From ac61598ad476e1c2d79af7da16323216d1153a5f Mon Sep 17 00:00:00 2001 From: iPoe Date: Sun, 12 Mar 2023 23:09:52 -0500 Subject: [PATCH] Add custom text for the in & out tags. --- supervision/detection/line_counter.py | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/supervision/detection/line_counter.py b/supervision/detection/line_counter.py index e40ead8e..5daccd59 100644 --- a/supervision/detection/line_counter.py +++ b/supervision/detection/line_counter.py @@ -6,6 +6,7 @@ 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 typing import Optional class LineZone: @@ -70,7 +71,6 @@ class LineZone: else: self.out_count += 1 - class LineZoneAnnotator: def __init__( self, @@ -81,6 +81,8 @@ class LineZoneAnnotator: text_scale: float = 0.5, text_offset: float = 1.5, text_padding: int = 10, + custom_in_text: Optional[str] = None, + custom_out_text: Optional[str] = None, ): """ Initialize the LineCounterAnnotator object with default values. @@ -102,6 +104,8 @@ class LineZoneAnnotator: self.text_scale: float = text_scale self.text_offset: float = text_offset self.text_padding: int = text_padding + self.custom_in_text: str = custom_in_text + self.custom_out_text: str = custom_out_text def annotate(self, frame: np.ndarray, line_counter: LineZone) -> np.ndarray: """ @@ -141,8 +145,9 @@ class LineZoneAnnotator: lineType=cv2.LINE_AA, ) - in_text = f"in: {line_counter.in_count}" - out_text = f"out: {line_counter.out_count}" + in_text = f"{self.custom_in_text}: {line_counter.in_count}" if self.custom_in_text is not None else f"in: {line_counter.in_count}" + out_text = f"{self.custom_out_text}: {line_counter.out_count}" if self.custom_out_text is not None else 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