Untitled
unknown
plain_text
a year ago
4.2 kB
7
Indexable
def annotate(
self,
frame: np.ndarray,
line_zones: List[LineZone],
line_zone_labels: Optional[List[str]] = None,
) -> np.ndarray:
"""
Draws a table with the number of objects of each class that crossed each line.
Attributes:
frame (np.ndarray): The image on which the table will be drawn.
line_zones (List[LineZone]): The line zones to be annotated.
line_zone_labels (Optional[List[str]]): The labels, one for each
line zone. If not provided, the default labels will be used.
Returns:
(np.ndarray): The image with the table drawn on it.
"""
if line_zone_labels is None:
line_zone_labels = [f"Lan {i + 1}:" for i in range(len(line_zones))]
if len(line_zones) != len(line_zone_labels):
raise ValueError("The number of line zones and their labels must match.")
text_lines = ["He thong dem so luong", "phuong tien luu thong:"]
for line_zone, line_zone_label in zip(line_zones, line_zone_labels):
text_lines.append(line_zone_label)
class_id_to_name = line_zone.class_id_to_name
for direction, count_per_class in [
("Vao", line_zone.in_count_per_class),
("Ra", line_zone.out_count_per_class),
]:
if not count_per_class:
continue
text_lines.append(f" {direction}:")
for class_id, count in count_per_class.items():
class_name = (
class_id_to_name.get(class_id, str(class_id))
if not self.force_draw_class_ids
else str(class_id)
)
text_lines.append(f" {class_name}: {count}")
text_lines.append("MQ ICT Solutions @ 2024")
table_width, table_height = 0, 0
for line in text_lines:
text_width, text_height = cv2.getTextSize(
line, cv2.FONT_HERSHEY_SIMPLEX, self.text_scale, self.text_thickness
)[0]
text_height += TEXT_MARGIN
table_width = max(table_width, text_width)
table_height += text_height
table_width += 2 * self.table_padding
table_height += 2 * self.table_padding
table_max_height = frame.shape[0] - 2 * self.table_margin
table_height = min(table_height, table_max_height)
table_width = min(table_width, self.table_max_width)
position_map = {
Position.TOP_LEFT: (self.table_margin, self.table_margin),
Position.TOP_RIGHT: (
frame.shape[1] - table_width - self.table_margin,
self.table_margin,
),
Position.BOTTOM_LEFT: (
self.table_margin,
frame.shape[0] - table_height - self.table_margin,
),
Position.BOTTOM_RIGHT: (
frame.shape[1] - table_width - self.table_margin,
frame.shape[0] - table_height - self.table_margin,
),
}
table_x1, table_y1 = position_map[self.table_position]
table_rect = Rect(
x=table_x1, y=table_y1, width=table_width, height=table_height
)
frame = draw_rectangle(
scene=frame, rect=table_rect, color=self.table_color, thickness=-1
)
for i, line in enumerate(text_lines):
_, text_height = cv2.getTextSize(
line, cv2.FONT_HERSHEY_SIMPLEX, self.text_scale, self.text_thickness
)[0]
text_height += TEXT_MARGIN
anchor_x = table_x1 + self.table_padding
anchor_y = table_y1 + self.table_padding + (i + 1) * text_height
cv2.putText(
img=frame,
text=line,
org=(anchor_x, anchor_y),
fontFace=cv2.FONT_HERSHEY_SIMPLEX,
fontScale=self.text_scale,
color=self.text_color.as_bgr(),
thickness=self.text_thickness,
lineType=cv2.LINE_AA,
)
return frameEditor is loading...
Leave a Comment