From 9210ccf8b75e28051503ae3cddcdaaf2caad9ebb Mon Sep 17 00:00:00 2001 From: Kader Miyanyedi Date: Mon, 29 Jan 2024 22:30:34 +0300 Subject: [PATCH] refactor(line_zone_annotator): use draw_text function in the annotate_count method and add center method to vector --- supervision/detection/line_counter.py | 75 +++++++++++---------------- supervision/draw/utils.py | 19 ++++--- supervision/geometry/core.py | 13 +++++ 3 files changed, 57 insertions(+), 50 deletions(-) diff --git a/supervision/detection/line_counter.py b/supervision/detection/line_counter.py index f2b7906e..b2d3b259 100644 --- a/supervision/detection/line_counter.py +++ b/supervision/detection/line_counter.py @@ -5,7 +5,8 @@ import numpy as np from supervision.detection.core import Detections from supervision.draw.color import Color -from supervision.geometry.core import Point, Position, Rect, Vector +from supervision.draw.utils import draw_text +from supervision.geometry.core import Point, Position, Vector class LineZone: @@ -209,7 +210,7 @@ class LineZoneAnnotator: frame: np.ndarray, line_counter: LineZone, text: str, - in_count: bool, + count_pos: bool, ) -> None: """This method is drawing the text on the frame. @@ -218,51 +219,32 @@ class LineZoneAnnotator: line_counter (LineCounter): The line counter that will be used to draw the line. text (str): The text that will be drawn. - in_count (bool): Whether to display the in count or not. + count_pos (bool): Whether to display the in count or not. """ - (text_width, text_height), _ = cv2.getTextSize( + text_width, text_height = cv2.getTextSize( text, cv2.FONT_HERSHEY_SIMPLEX, self.text_scale, self.text_thickness - ) + )[0] - text_x = ( - line_counter.vector.end.x + line_counter.vector.start.x - text_width - ) / 2 + center_text_anchor = Vector( + start=line_counter.vector.start, end=line_counter.vector.end + ).center - text_y = ( - line_counter.vector.end.y + line_counter.vector.start.y + text_height - ) / 2 - - if in_count: - # if the text is an in count, we want to display it above the line - text_y -= self.text_offset * text_height + if count_pos: + center_text_anchor.y -= int(self.text_offset * text_height) else: - # if the text is an out count, we want to display it below the line - text_y += self.text_offset * text_height + center_text_anchor.y += int(self.text_offset * text_height) - text_background_rect = Rect( - x=text_x, - y=text_y - text_height, - width=text_width, - height=text_height, - ).pad(padding=self.text_padding) - - cv2.rectangle( - frame, - text_background_rect.top_left.as_xy_int_tuple(), - text_background_rect.bottom_right.as_xy_int_tuple(), - self.color.as_bgr(), - -1, - ) - - cv2.putText( - frame, - text, - (int(text_x), int(text_y)), - cv2.FONT_HERSHEY_SIMPLEX, - self.text_scale, - self.text_color.as_bgr(), - self.text_thickness, - cv2.LINE_AA, + draw_text( + scene=frame, + text=text, + text_anchor=center_text_anchor, + text_color=self.text_color, + text_scale=self.text_scale, + text_thickness=self.text_thickness, + text_padding=self.text_padding, + background_color=self.color, + text_width=text_width, + text_height=text_height, ) def annotate(self, frame: np.ndarray, line_counter: LineZone) -> np.ndarray: @@ -311,7 +293,10 @@ class LineZoneAnnotator: else f"in: {line_counter.in_count}" ) self._annotate_count( - line_counter=line_counter, text=in_text, in_count=True, frame=frame + frame=frame, + line_counter=line_counter, + text=in_text, + count_pos=True, ) if self.display_out_count: @@ -321,7 +306,9 @@ class LineZoneAnnotator: else f"out: {line_counter.out_count}" ) self._annotate_count( - line_counter=line_counter, text=out_text, in_count=False, frame=frame + frame=frame, + line_counter=line_counter, + text=out_text, + count_pos=False, ) - return frame diff --git a/supervision/draw/utils.py b/supervision/draw/utils.py index a6e600e3..982d9f1f 100644 --- a/supervision/draw/utils.py +++ b/supervision/draw/utils.py @@ -111,6 +111,8 @@ def draw_text( text_padding: int = 10, text_font: int = cv2.FONT_HERSHEY_SIMPLEX, background_color: Optional[Color] = None, + text_width: Optional[int] = None, + text_height: Optional[int] = None, ) -> np.ndarray: """ Draw text with background on a scene. @@ -129,6 +131,10 @@ def draw_text( Defaults to cv2.FONT_HERSHEY_SIMPLEX. background_color (Color, optional): The color of the background rectangle, if one is to be drawn. Defaults to None. + text_width (Optional[int], optional): The width of the text, if known. + Defaults to None. If None, the width will be calculated. + text_height (Optional[int], optional): The height of the text, if known. + Defaults to None. If None, the height will be calculated. Returns: np.ndarray: The input scene with the text drawn on it. @@ -142,12 +148,13 @@ def draw_text( scene = draw_text(scene=scene, text="Hello, world!",text_anchor=text_anchor) ``` """ - text_width, text_height = cv2.getTextSize( - text=text, - fontFace=text_font, - fontScale=text_scale, - thickness=text_thickness, - )[0] + if text_width is None or text_height is None: + text_width, text_height = cv2.getTextSize( + text=text, + fontFace=text_font, + fontScale=text_scale, + thickness=text_thickness, + )[0] text_rect = Rect( x=text_anchor.x - text_width // 2, y=text_anchor.y - text_height // 2, diff --git a/supervision/geometry/core.py b/supervision/geometry/core.py index 3e4ee462..a87ca949 100644 --- a/supervision/geometry/core.py +++ b/supervision/geometry/core.py @@ -56,6 +56,19 @@ class Vector: dy = self.end.y - self.start.y return sqrt(dx**2 + dy**2) + @property + def center(self) -> Point: + """ + Calculate the center point of the vector. + + Returns: + Point: The center point of the vector. + """ + return Point( + x=(self.start.x + self.end.x) // 2, + y=(self.start.y + self.end.y) // 2, + ) + def cross_product(self, point: Point) -> float: """ Calculate the 2D cross product (also known as the vector product or outer