scripts/plot_tracks.py

import supervision as sv
import cv2
import os


def annotate_image(image, detections, class_names, trace_annotator):
    """
    Draws bounding boxes, class labels with track IDs, and trace lines on the image.

    Args:
        image (np.ndarray): The image to annotate.
        detections (sv.Detections): Detection results.
        class_names (dict): Class ID to name mapping.
        trace_annotator (sv.TraceAnnotator): Annotator for trace lines.

    Returns:
        np.ndarray: Annotated image.
    """
    if len(detections) > 0:
        # Initialize annotators, trace annotator should be initialized only once
        # Draw box
        _, width = image.shape[:2]
        text_size = 2.0*width/1024
        box_thickness = int(6 * width / 1024)  # scale box thickness
        text_thickness = max(1, int(6 * width / 1024))  # scale text thickness, min 1

        box_annotator = sv.BoxAnnotator(thickness=box_thickness)
        label_annotator = sv.LabelAnnotator(text_scale=text_size, text_thickness=text_thickness)

        # Build labels with track IDs
        labels = []
        for i in range(len(detections)):
            class_id = detections.class_id[i]
            conf = detections.confidence[i]
            track_id = detections.tracker_id[i] if detections.tracker_id is not None else None

            if track_id is not None:
                label = f"ID: {track_id} {class_names[class_id]} {conf:.2f}"
            else:
                label = f"{class_names[class_id]} {conf:.2f}"
            labels.append(label)

        # Apply all annotators
        image = box_annotator.annotate(scene=image, detections=detections)
        image = label_annotator.annotate(scene=image, detections=detections, labels=labels)
        image = trace_annotator.annotate(scene=image, detections=detections)

    return image


def process_video_frame(frame, detections, class_names, trace_annotator, plot_dets=True, show=False, video_writer=None):
    """
    Handles each frame of video for annotation, display, and saving.

    Args:
        frame (np.ndarray): Video frame.
        detections (sv.Detections): Detection results.
        class_names (dict): Class ID to name mapping.
        trace_annotator (sv.TraceAnnotator): Annotator for trace lines.
        plot_dets (bool): Save annotated video frame.
        show (bool): Display annotated frame in a window.
    """
    annotated_frame = annotate_image(frame, detections, class_names, trace_annotator)

    if show:
        h, w = frame.shape[:2]
        scale = min(750 / w, 750 / h)
        resized = cv2.resize(annotated_frame, (int(w * scale), int(h * scale)))
        cv2.imshow("Detection", resized)
        cv2.waitKey(1)

    return annotated_frame