xiaoxue_admin
/
supervision
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Merge pull request #854 from jcruz-ferreyra/lc/diagonal_support 

LineZoneAnnotator: Align text to line counter in non-horizontal lines
This commit is contained in:
LinasKo 2024-09-26 13:01:42 +03:00 committed by GitHub
parent d4eb957201 c81aec9d1b
commit e63611deb3
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 333 additions and 84 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

406
supervision/detection/line_zone.py
View File

@ -1,5 +1,7 @@
import math
import warnings
from typing import Dict, Iterable, Optional, Tuple
from functools import lru_cache
from typing import Any, Dict, Iterable, Optional, Tuple
import cv2
import numpy as np
@ -9,6 +11,7 @@ from supervision.detection.utils import cross_product
from supervision.draw.color import Color
from supervision.draw.utils import draw_text
from supervision.geometry.core import Point, Position, Vector
from supervision.utils.image import overlay_image
from supervision.utils.internal import SupervisionWarnings
@ -199,9 +202,9 @@ class LineZone:
class LineZoneAnnotator:
def __init__(
self,
thickness: float = 2,
thickness: int = 2,
color: Color = Color.WHITE,
text_thickness: float = 2,
text_thickness: int = 2,
text_color: Color = Color.BLACK,
text_scale: float = 0.5,
text_offset: float = 1.5,
@ -210,86 +213,67 @@ class LineZoneAnnotator:
custom_out_text: Optional[str] = None,
display_in_count: bool = True,
display_out_count: bool = True,
display_text_box: bool = True,
text_orient_to_line: bool = False,
text_centered: bool = True,
):
"""
Initialize the LineCounterAnnotator object with default values.
A class for drawing the `LineZone` and its detected object count
on an image.
Attributes:
thickness (float): The thickness of the line that will be drawn.
color (Color): The color of the line that will be drawn.
text_thickness (float): The thickness of the text that will be drawn.
text_color (Color): The color of the text that will be drawn.
text_scale (float): The scale of the text that will be drawn.
text_offset (float): The offset of the text that will be drawn.
text_padding (int): The padding of the text that will be drawn.
display_in_count (bool): Whether to display the in count or not.
display_out_count (bool): Whether to display the out count or not.
thickness (int): Line thickness.
color (Color): Line color.
text_thickness (int): Text thickness.
text_color (Color): Text color.
text_scale (float): Text scale.
text_offset (float): How far the text will be from the line.
text_padding (int): The empty space in the text box, surrounding the text.
custom_in_text (Optional[str]): Write something else instead of "in".
custom_out_text (Optional[str]): Write something else instead of "out".
display_in_count (bool): Pass `False` to hide the "in" count.
display_out_count (bool): Pass `False` to hide the "out" count.
display_text_box (bool): Pass `False` to hide the text background box.
text_orient_to_line (bool): ⭐ Match text orientation to the line.
Recommended to set to `True`.
text_centered (bool): Pass `False` to disable text centering. Useful
when the label overlaps something important.
"""
self.thickness: float = thickness
self.thickness: int = thickness
self.color: Color = color
self.text_thickness: float = text_thickness
self.text_thickness: int = 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
self.custom_in_text: str = custom_in_text
self.custom_out_text: str = custom_out_text
self.in_text: str = custom_in_text if custom_in_text else "in"
self.out_text: str = custom_out_text if custom_out_text else "out"
self.display_in_count: bool = display_in_count
self.display_out_count: bool = display_out_count
def _annotate_count(
self,
frame: np.ndarray,
center_text_anchor: Point,
text: str,
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.
center_text_anchor: The center point that the text will be drawn.
text (str): The text that will be drawn.
is_in_count (bool): Whether to display the in count or out count.
"""
_, text_height = cv2.getTextSize(
text, cv2.FONT_HERSHEY_SIMPLEX, self.text_scale, self.text_thickness
)[0]
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)
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,
)
self.display_text_box: bool = display_text_box
self.text_orient_to_line: bool = text_orient_to_line
self.text_centered: bool = text_centered
def annotate(self, frame: np.ndarray, line_counter: LineZone) -> np.ndarray:
"""
Draws the line on the frame using the line_counter provided.
Draws the line on the frame using the line zone provided.
Attributes:
frame (np.ndarray): The image on which the line will be drawn.
line_counter (LineCounter): The line counter
line_counter (LineZone): The line zone
that will be used to draw the line.
Returns:
np.ndarray: The image with the line drawn on it.
"""
line_start = line_counter.vector.start.as_xy_int_tuple()
line_end = line_counter.vector.end.as_xy_int_tuple()
cv2.line(
frame,
line_counter.vector.start.as_xy_int_tuple(),
line_counter.vector.end.as_xy_int_tuple(),
line_start,
line_end,
self.color.as_bgr(),
self.thickness,
lineType=cv2.LINE_AA,
@ -297,7 +281,7 @@ class LineZoneAnnotator:
)
cv2.circle(
frame,
line_counter.vector.start.as_xy_int_tuple(),
line_start,
radius=5,
color=self.text_color.as_bgr(),
thickness=-1,
@ -305,40 +289,294 @@ class LineZoneAnnotator:
)
cv2.circle(
frame,
line_counter.vector.end.as_xy_int_tuple(),
line_end,
radius=5,
color=self.text_color.as_bgr(),
thickness=-1,
lineType=cv2.LINE_AA,
)
text_anchor = Vector(
start=line_counter.vector.start, end=line_counter.vector.end
in_text = f"{self.in_text}: {line_counter.in_count}"
out_text = f"{self.out_text}: {line_counter.out_count}"
line_angle_degrees = self._get_line_angle(line_counter)
for text, is_shown, is_in_count in [
(in_text, self.display_in_count, True),
(out_text, self.display_out_count, False),
]:
if not is_shown:
continue
if line_angle_degrees == 0 or not self.text_orient_to_line:
self._draw_basic_label(
frame=frame,
line_center=line_counter.vector.center,
text=text,
is_in_count=is_in_count,
)
else:
self._draw_oriented_label(
frame=frame,
line_zone=line_counter,
text=text,
is_in_count=is_in_count,
)
return frame
def _get_line_angle(self, line_zone: LineZone) -> float:
"""
Calculate the line counter angle (in degrees).
Args:
line_zone (LineZone): The line zone object.
Returns:
float: Line counter angle, in degrees.
"""
start_point = line_zone.vector.start.as_xy_int_tuple()
end_point = line_zone.vector.end.as_xy_int_tuple()
delta_x = end_point[0] - start_point[0]
delta_y = end_point[1] - start_point[1]
if delta_x == 0:
line_angle = 90.0
line_angle += 180 if delta_y < 0 else 0
else:
line_angle = math.degrees(math.atan(delta_y / delta_x))
line_angle += 180 if delta_x < 0 else 0
return line_angle
def _calculate_anchor_in_frame(
self,
line_zone: LineZone,
text_width: int,
text_height: int,
is_in_count: bool,
label_dimension: int,
) -> Tuple[int, int]:
"""
Calculate insertion anchor in frame to position the center of the count image.
Args:
line_zone (LineZone): The line counter object used for counting.
text_width (int): Text width.
text_height (int): Text height.
is_in_count (bool): Whether the count should be placed over or below line.
label_dimension (int): Size of the label image. Assumes the
label is rectangular.
Returns:
Tuple[int, int]: xy, pont in an image where the label will be placed.
"""
line_angle = self._get_line_angle(line_zone)
if self.text_centered:
mid_point = Vector(
start=line_zone.vector.start, end=line_zone.vector.end
).center.as_xy_int_tuple()
anchor = list(mid_point)
else:
end_point = line_zone.vector.end.as_xy_int_tuple()
anchor = list(end_point)
move_along_x = int(
math.cos(math.radians(line_angle))
* (text_width / 2 + self.text_padding)
)
move_along_y = int(
math.sin(math.radians(line_angle))
* (text_width / 2 + self.text_padding)
)
anchor[0] -= move_along_x
anchor[1] -= move_along_y
move_perpendicular_x = int(
math.sin(math.radians(line_angle)) * (self.text_offset * text_height)
)
move_perpendicular_y = int(
math.cos(math.radians(line_angle)) * (self.text_offset * text_height)
)
if self.display_in_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}"
)
self._annotate_count(
frame=frame,
center_text_anchor=text_anchor.center,
text=in_text,
is_in_count=True,
)
if is_in_count:
anchor[0] += move_perpendicular_x
anchor[1] -= move_perpendicular_y
else:
anchor[0] -= move_perpendicular_x
anchor[1] += move_perpendicular_y
x1 = max(anchor[0] - label_dimension // 2, 0)
y1 = max(anchor[1] - label_dimension // 2, 0)
return x1, y1
def _draw_basic_label(
self,
frame: np.ndarray,
line_center: Point,
text: str,
is_in_count: bool,
) -> np.ndarray:
"""
Draw the count label on the frame. For example: "out: 7".
The label contains horizontal text and is not rotated.
Args:
frame (np.ndarray): The entire scene, on which the label will be placed.
line_center (Point): The center of the line zone.
text (str): The text that will be drawn.
is_in_count (bool): Whether to display the in count (above line)
or out count (below line).
Returns:
np.ndarray: The scene with the label drawn on it.
"""
_, text_height = cv2.getTextSize(
text, cv2.FONT_HERSHEY_SIMPLEX, self.text_scale, self.text_thickness
)[0]
if is_in_count:
line_center.y -= int(self.text_offset * text_height)
else:
line_center.y += int(self.text_offset * text_height)
draw_text(
scene=frame,
text=text,
text_anchor=line_center,
text_color=self.text_color,
text_scale=self.text_scale,
text_thickness=self.text_thickness,
text_padding=self.text_padding,
background_color=self.color if self.display_text_box else None,
)
if self.display_out_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}"
)
self._annotate_count(
frame=frame,
center_text_anchor=text_anchor.center,
text=out_text,
is_in_count=False,
)
return frame
def _draw_oriented_label(
self,
frame: np.ndarray,
line_zone: LineZone,
text: str,
is_in_count: bool,
) -> np.ndarray:
"""
Draw the count label on the frame. For example: "out: 7".
The label is oriented to match the line angle.
Args:
frame (np.ndarray): The entire scene, on which the label will be placed.
line_zone (LineZone): The line zone responsible for counting
objects crossing it.
text (str): The text that will be drawn.
is_in_count (bool): Whether to display the in count (above line)
or out count (below line).
Returns:
np.ndarray: The scene with the label drawn on it.
"""
line_angle_degrees = self._get_line_angle(line_zone)
label_image = self._make_label_image(
text,
text_scale=self.text_scale,
text_thickness=self.text_thickness,
text_padding=self.text_padding,
text_color=self.text_color,
text_box_show=self.display_text_box,
text_box_color=self.color,
line_angle_degrees=line_angle_degrees,
)
assert label_image.shape[0] == label_image.shape[1]
text_width, text_height = cv2.getTextSize(
text, cv2.FONT_HERSHEY_SIMPLEX, self.text_scale, self.text_thickness
)[0]
label_anchor = self._calculate_anchor_in_frame(
line_zone=line_zone,
text_width=text_width,
text_height=text_height,
is_in_count=is_in_count,
label_dimension=label_image.shape[0],
)
frame = overlay_image(frame, label_image, label_anchor)
return frame
@staticmethod
@lru_cache(maxsize=32)
def _make_label_image(
text: str,
*,
text_scale: float,
text_thickness: int,
text_padding: int,
text_color: Color,
text_box_show: bool,
text_box_color: Color,
line_angle_degrees: float,
) -> np.ndarray:
"""
Create the small text box displaying line zone count. E.g. "out: 7".
Args:
text (str): The text to display.
text_scale (float): The scale of the text.
text_thickness (int): The thickness of the text.
text_padding (int): The padding around the text.
text_color (Color): The color of the text.
text_box_show (bool): Whether to display the text box.
text_box_color (Color): The color of the text box.
line_angle_degrees (float): The angle of the line in degrees.
Returns:
np.ndarray: The label of shape (H, W, 4), in BGRA format.
"""
text_width, text_height = cv2.getTextSize(
text, cv2.FONT_HERSHEY_SIMPLEX, text_scale, text_thickness
)[0]
annotation_dim = int((max(text_width, text_height) + text_padding * 2) * 1.5)
annotation_shape = (annotation_dim, annotation_dim)
annotation_center = Point(annotation_dim // 2, annotation_dim // 2)
annotation = np.zeros((*annotation_shape, 3), dtype=np.uint8)
annotation_alpha = np.zeros((*annotation_shape, 1), dtype=np.uint8)
text_args: Dict[str, Any] = dict(
text=text,
text_anchor=annotation_center,
text_scale=text_scale,
text_thickness=text_thickness,
text_padding=text_padding,
)
draw_text(
scene=annotation,
text_color=text_color,
background_color=text_box_color if text_box_show else None,
**text_args,
)
draw_text(
scene=annotation_alpha,
text_color=Color.WHITE,
background_color=Color.WHITE if text_box_show else None,
**text_args,
)
annotation = np.dstack((annotation, annotation_alpha))
# Make sure text is displayed upright
if 90 < line_angle_degrees % 360 < 270:
annotation = cv2.flip(annotation, flipCode=-1).astype(np.uint8)
rotation_angle = -line_angle_degrees
rotation_matrix = cv2.getRotationMatrix2D(
annotation_center.as_xy_float_tuple(), rotation_angle, scale=1
)
annotation = cv2.warpAffine(annotation, rotation_matrix, annotation_shape)
return annotation

11
supervision/draw/color.py
View File

@ -255,6 +255,17 @@ class Color:
def ROBOFLOW(cls) -> Color:
return Color.from_hex("#A351FB")
def __hash__(self):
return hash((self.r, self.g, self.b))
def __eq__(self, other):
return (
isinstance(other, Color)
and self.r == other.r
and self.g == other.g
and self.b == other.b
)
@dataclass
class ColorPalette: