From e0294cd296ca39268a8ce3989a85ecaa8cd5eda0 Mon Sep 17 00:00:00 2001 From: Kader Miyanyedi Date: Tue, 30 Jan 2024 14:28:15 +0300 Subject: [PATCH] refactor(line_zone_annotator): refactor _annotate_count method --- supervision/detection/line_counter.py | 31 ++++++++++++--------------- supervision/draw/utils.py | 28 +++++++++++------------- supervision/geometry/core.py | 4 ++-- 3 files changed, 28 insertions(+), 35 deletions(-) diff --git a/supervision/detection/line_counter.py b/supervision/detection/line_counter.py index b2d3b259..43720e27 100644 --- a/supervision/detection/line_counter.py +++ b/supervision/detection/line_counter.py @@ -208,28 +208,23 @@ class LineZoneAnnotator: def _annotate_count( self, frame: np.ndarray, - line_counter: LineZone, + center_text_anchor: Point, text: str, - count_pos: bool, + is_in_count: bool, ) -> None: """This method is drawing the text on the frame. Args: frame (np.ndarray): The image on which the text will be drawn. - line_counter (LineCounter): The line counter - that will be used to draw the line. + center_text_anchor: The center point that the text will be drawn. text (str): The text that will be drawn. - count_pos (bool): Whether to display the in count or not. + is_in_count (bool): Whether to display the in count or not. """ - text_width, text_height = cv2.getTextSize( + _, text_height = cv2.getTextSize( text, cv2.FONT_HERSHEY_SIMPLEX, self.text_scale, self.text_thickness )[0] - center_text_anchor = Vector( - start=line_counter.vector.start, end=line_counter.vector.end - ).center - - if count_pos: + if is_in_count: center_text_anchor.y -= int(self.text_offset * text_height) else: center_text_anchor.y += int(self.text_offset * text_height) @@ -243,8 +238,6 @@ class LineZoneAnnotator: 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: @@ -286,6 +279,10 @@ class LineZoneAnnotator: lineType=cv2.LINE_AA, ) + text_anchor = Vector( + start=line_counter.vector.start, end=line_counter.vector.end + ) + if self.display_in_count: in_text = ( f"{self.custom_in_text}: {line_counter.in_count}" @@ -294,9 +291,9 @@ class LineZoneAnnotator: ) self._annotate_count( frame=frame, - line_counter=line_counter, + center_text_anchor=text_anchor.center, text=in_text, - count_pos=True, + is_in_count=True, ) if self.display_out_count: @@ -307,8 +304,8 @@ class LineZoneAnnotator: ) self._annotate_count( frame=frame, - line_counter=line_counter, + center_text_anchor=text_anchor.center, text=out_text, - count_pos=False, + is_in_count=False, ) return frame diff --git a/supervision/draw/utils.py b/supervision/draw/utils.py index 982d9f1f..7f82b008 100644 --- a/supervision/draw/utils.py +++ b/supervision/draw/utils.py @@ -111,8 +111,6 @@ 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. @@ -131,10 +129,6 @@ 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. @@ -148,16 +142,18 @@ def draw_text( scene = draw_text(scene=scene, text="Hello, world!",text_anchor=text_anchor) ``` """ - 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_width, text_height = cv2.getTextSize( + text=text, + fontFace=text_font, + fontScale=text_scale, + thickness=text_thickness, + )[0] + + text_anchor_x, text_anchor_y = text_anchor.as_xy_int_tuple() + text_rect = Rect( - x=text_anchor.x - text_width // 2, - y=text_anchor.y - text_height // 2, + x=text_anchor_x - text_width // 2, + y=text_anchor_y - text_height // 2, width=text_width, height=text_height, ).pad(text_padding) @@ -170,7 +166,7 @@ def draw_text( cv2.putText( img=scene, text=text, - org=(text_anchor.x - text_width // 2, text_anchor.y + text_height // 2), + org=(text_anchor_x - text_width // 2, text_anchor_y + text_height // 2), fontFace=text_font, fontScale=text_scale, color=text_color.as_bgr(), diff --git a/supervision/geometry/core.py b/supervision/geometry/core.py index a87ca949..39d42c60 100644 --- a/supervision/geometry/core.py +++ b/supervision/geometry/core.py @@ -65,8 +65,8 @@ class Vector: 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, + x=(self.start.x + self.end.x) / 2, + y=(self.start.y + self.end.y) / 2, ) def cross_product(self, point: Point) -> float: