src/visualization/annotator.py

"""OpenCV annotation utilities."""

from __future__ import annotations

import hashlib
from collections.abc import Sequence

import cv2
import numpy as np

from src.tracking.tracker import TrackHistory, TrackedObject
from src.visualization.trajectory import draw_trajectory


class TrackAnnotator:
    """Draw bounding boxes, labels, confidences, and trajectory trails."""

    def __init__(
        self,
        draw_trajectories: bool = True,
        trajectory_length: int = 40,
        line_thickness: int = 2,
    ) -> None:
        self.draw_trajectories = draw_trajectories
        self.trajectory_length = trajectory_length
        self.line_thickness = line_thickness

    @staticmethod
    def _color_for_id(track_id: int) -> tuple[int, int, int]:
        digest = hashlib.md5(str(track_id).encode("utf-8")).digest()
        return int(digest[0]), int(digest[1]), int(digest[2])

    def draw(
        self,
        frame: np.ndarray,
        tracks: Sequence[TrackedObject],
        history: TrackHistory,
    ) -> np.ndarray:
        annotated = frame.copy()
        for track in tracks:
            color = self._color_for_id(track.id)
            x1, y1, x2, y2 = (int(v) for v in track.bbox)
            cv2.rectangle(annotated, (x1, y1), (x2, y2), color, self.line_thickness)

            label = f"ID: {track.id} | {track.confidence:.2f}"
            (text_width, text_height), baseline = cv2.getTextSize(
                label, cv2.FONT_HERSHEY_SIMPLEX, 0.55, 2
            )
            label_y = max(y1 - text_height - baseline - 4, 0)
            cv2.rectangle(
                annotated,
                (x1, label_y),
                (x1 + text_width + 8, label_y + text_height + baseline + 6),
                color,
                -1,
            )
            cv2.putText(
                annotated,
                label,
                (x1 + 4, label_y + text_height + 2),
                cv2.FONT_HERSHEY_SIMPLEX,
                0.55,
                (255, 255, 255),
                2,
                cv2.LINE_AA,
            )

            if self.draw_trajectories:
                points = history.get_recent_points(track.id, self.trajectory_length)
                draw_trajectory(annotated, points, color, max(1, self.line_thickness - 1))

        return annotated