agents/annotator.py

"""
Draw bounding boxes, labels, trajectories, and annotations on camera frames.

Used by all agents to produce annotated_frame_base64 output.
Requires OpenCV (cv2). Falls back gracefully if not available.
"""

import base64
from typing import Optional

try:
    import cv2
    import numpy as np
    HAS_CV2 = True
except ImportError:
    HAS_CV2 = False

# Color scheme per agent type (BGR for OpenCV)
COLORS = {
    "detect": (0, 255, 0),       # Green bboxes
    "alert": (0, 0, 255),        # Red for alerts
    "count": (0, 165, 255),      # Orange numbered markers
    "track": (255, 0, 255),      # Magenta trajectories
    "ocr": (255, 255, 0),        # Cyan text regions
    "caption": (200, 200, 200),  # Gray labels
    "reason": (0, 255, 255),     # Yellow areas of concern
    "default": (0, 255, 0),
}


class FrameAnnotator:
    """Draw annotations on camera frames."""

    @staticmethod
    def annotate(frame, detections: list[dict] = None,
                 agent_type: str = "detect",
                 counts: dict = None,
                 text_regions: list[dict] = None,
                 tracks: list[dict] = None,
                 alert: dict = None) -> Optional[object]:
        """Draw all annotations on frame. Returns annotated copy or None."""
        if not HAS_CV2 or frame is None:
            return None

        annotated = frame.copy()
        h, w = annotated.shape[:2]
        color = COLORS.get(agent_type, COLORS["default"])

        # Draw bounding boxes + labels
        if detections:
            for i, det in enumerate(detections):
                bbox = det.get("bbox", [])
                if len(bbox) == 4:
                    x1 = int(bbox[0] * w)
                    y1 = int(bbox[1] * h)
                    x2 = int(bbox[2] * w)
                    y2 = int(bbox[3] * h)
                    cv2.rectangle(annotated, (x1, y1), (x2, y2), color, 2)

                    label = det.get("label", "")
                    conf = det.get("confidence", 0)
                    count_idx = det.get("count_index")

                    # Count mode: numbered circle
                    if count_idx:
                        cx, cy = (x1 + x2) // 2, max(y1 - 15, 15)
                        cv2.circle(annotated, (cx, cy), 14, color, -1)
                        cv2.putText(annotated, str(count_idx), (cx - 7, cy + 5),
                                    cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 0), 2)

                    if label:
                        text = f"{label} {conf:.0%}" if conf else label
                        # Background for text
                        (tw, th), _ = cv2.getTextSize(text, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 1)
                        cv2.rectangle(annotated, (x1, y1 - th - 8), (x1 + tw + 4, y1), color, -1)
                        cv2.putText(annotated, text, (x1 + 2, y1 - 4),
                                    cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 0), 1)

        # Draw text regions (OCR)
        if text_regions:
            for tr in text_regions:
                bbox = tr.get("bbox", [])
                if len(bbox) == 4:
                    x1, y1 = int(bbox[0] * w), int(bbox[1] * h)
                    x2, y2 = int(bbox[2] * w), int(bbox[3] * h)
                    cv2.rectangle(annotated, (x1, y1), (x2, y2), COLORS["ocr"], 2)
                    text = tr.get("text", "")
                    if text:
                        cv2.putText(annotated, text, (x1, y1 - 5),
                                    cv2.FONT_HERSHEY_SIMPLEX, 0.5, COLORS["ocr"], 1)

        # Draw trajectories (Track)
        if tracks:
            for track in tracks:
                trajectory = track.get("trajectory", [])
                if len(trajectory) >= 2:
                    pts = [(int(p[0] * w), int(p[1] * h)) for p in trajectory]
                    for j in range(len(pts) - 1):
                        cv2.line(annotated, pts[j], pts[j + 1], COLORS["track"], 3)
                    if len(pts) >= 2:
                        cv2.arrowedLine(annotated, pts[-2], pts[-1], COLORS["track"], 3)
                    label = track.get("label", "")
                    oid = track.get("object_id", "")
                    if label and pts:
                        cv2.putText(annotated, f"{label} #{oid}", (pts[0][0], pts[0][1] - 10),
                                    cv2.FONT_HERSHEY_SIMPLEX, 0.5, COLORS["track"], 1)

        # Alert banner
        if alert and alert.get("severity") in ("HIGH", "CRITICAL"):
            banner_color = (0, 0, 200) if alert["severity"] == "CRITICAL" else (0, 100, 255)
            cv2.rectangle(annotated, (0, 0), (w, 35), banner_color, -1)
            desc = alert.get("description", alert.get("category", ""))[:80]
            cv2.putText(annotated, f"ALERT [{alert['severity']}]: {desc}",
                        (10, 25), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 255, 255), 2)

        # Count summary bar
        if counts and any(v for k, v in counts.items() if k != "total"):
            total = counts.get("total", sum(v for k, v in counts.items() if k != "total"))
            parts = [f"{v} {k}" for k, v in counts.items() if k != "total" and v > 0]
            text = " | ".join(parts) + f" (total: {total})"
            cv2.rectangle(annotated, (0, h - 35), (w, h), (0, 0, 0), -1)
            cv2.putText(annotated, text, (10, h - 10),
                        cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 1)

        return annotated

    @staticmethod
    def to_base64(frame, quality: int = 85) -> Optional[str]:
        """Encode frame as base64 JPEG."""
        if not HAS_CV2 or frame is None:
            return None
        _, jpeg = cv2.imencode(".jpg", frame, [cv2.IMWRITE_JPEG_QUALITY, quality])
        return base64.b64encode(jpeg.tobytes()).decode("utf-8")

    @staticmethod
    def annotate_and_encode(frame, quality: int = 85, **kwargs) -> Optional[str]:
        """Annotate frame and return as base64 JPEG. Returns None if cv2 unavailable."""
        annotated = FrameAnnotator.annotate(frame, **kwargs)
        if annotated is None:
            return None
        return FrameAnnotator.to_base64(annotated, quality)