Untitled
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 frame
Leave a Comment