Untitled

 avatar
unknown
plain_text
2 months ago
4.2 kB
3
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 frame
Leave a Comment