annotations.py

"""
annotations.py — Annotation drawing engine for Gridlock.
=========================================================
Extracted from run_inference.py and enhanced with support for:
  - Illegal parking annotations
  - Per-violation type badges
  - Confidence scores
  - Reusable from both CLI and API contexts
"""

import cv2
import numpy as np

from config import (
    CLR_BIKE, CLR_CAR, CLR_RIDER, CLR_NO_HELMET, CLR_PLATE,
    CLR_VIOLATION, CLR_OK, CLR_WRONGSIDE, CLR_SEATBELT, CLR_PARKING,
    HEAD_CROP_FRACTION, HEAD_CROP_MIN_PX,
)

FONT = cv2.FONT_HERSHEY_SIMPLEX


def put_label(img, text, x, y, color, bg_color=None, scale=0.55, thickness=1):
    """Draw text with an optional filled background pill."""
    (tw, th), baseline = cv2.getTextSize(text, FONT, scale, thickness)
    if bg_color is not None:
        pad = 4
        cv2.rectangle(
            img,
            (x - pad, y - th - pad),
            (x + tw + pad, y + baseline + pad),
            bg_color, -1,
        )
    cv2.putText(img, text, (x, y), FONT, scale, color, thickness, cv2.LINE_AA)


def draw_rounded_rect(img, x1, y1, x2, y2, color, radius=8, thickness=2):
    """Draw a rounded rectangle."""
    cv2.rectangle(img, (x1 + radius, y1), (x2 - radius, y2), color, thickness)
    cv2.rectangle(img, (x1, y1 + radius), (x2, y2 - radius), color, thickness)
    cv2.ellipse(img, (x1 + radius, y1 + radius), (radius, radius), 180, 0, 90, color, thickness)
    cv2.ellipse(img, (x2 - radius, y1 + radius), (radius, radius), 270, 0, 90, color, thickness)
    cv2.ellipse(img, (x1 + radius, y2 - radius), (radius, radius), 90, 0, 90, color, thickness)
    cv2.ellipse(img, (x2 - radius, y2 - radius), (radius, radius), 0, 0, 90, color, thickness)


def annotate_from_pipeline_result(img: np.ndarray, pipeline_result: dict) -> np.ndarray:
    """
    Draw full annotations on a copy of the image using the structured
    output from ParallelDetectionPipeline.process(return_annotations=True).

    Args:
        img: Original BGR image (numpy array).
        pipeline_result: Dict returned by pipeline.process() with
                         return_annotations=True.

    Returns:
        Annotated image (numpy array) with banner.
    """
    out = img.copy()
    h_img, w_img = img.shape[:2]

    annot = pipeline_result.get("annotation_data", {})
    vehicles = annot.get("vehicles", [])

    total_violations = 0
    total_bikes = pipeline_result.get("vehicles_detected", {}).get("bikes", 0)
    total_cars = pipeline_result.get("vehicles_detected", {}).get("cars", 0)

    # ── Draw vehicles ────────────────────────────────────────────────────────
    for v in vehicles:
        is_bike = v["is_bike"]
        x1, y1, x2, y2 = v["box"]
        is_violation = v["is_violation"]
        is_wrong_side = v["is_wrong_side"]
        num_riders = v["num_riders"]
        with_h = v["with_helmet"]
        without_h = v["without_helmet"]
        no_seatbelt = v["no_seatbelt"]
        plate_text = v["plate_text"]
        plate_box_abs = v["plate_box_abs"]
        rider_boxes = v["rider_boxes"]
        violation_types = v.get("violation_types", [])

        if is_violation:
            total_violations += 1

        # Vehicle bounding box
        vehicle_color = CLR_VIOLATION if is_violation else (CLR_BIKE if is_bike else CLR_CAR)
        cv2.rectangle(out, (x1, y1), (x2, y2), vehicle_color, 2)

        # Violation badge above vehicle
        badge_lines = []
        if is_violation:
            badge_lines.append(f"VIOLATION #{total_violations}")
            # Show specific violation types
            type_str = " | ".join(t.upper().replace("_", " ") for t in violation_types)
            if type_str:
                badge_lines.append(type_str)

        if is_wrong_side:
            badge_lines.append("WRONG SIDE!")

        if is_bike:
            badge_lines.append(f"Riders: {num_riders} Helmet OK: {with_h} No Helmet: {without_h}")
        else:
            if no_seatbelt > 0:
                badge_lines.append(f"No Seatbelt: {no_seatbelt}")

        if plate_text != "UNKNOWN":
            badge_lines.append(f"Plate: {plate_text}")

        badge_y = max(y1 - 6 - 16 * len(badge_lines), 5)
        for li, line in enumerate(badge_lines):
            bg = (0, 0, 180) if (li == 0 and is_violation) else (30, 30, 30)
            txt_color = (255, 255, 255)
            put_label(out, line, x1, badge_y + li * 16, txt_color, bg, scale=0.35, thickness=1)

        # Draw rider boxes (bikes only)
        if is_bike and rider_boxes:
            for p_box in rider_boxes:
                px1, py1, px2, py2 = map(int, p_box)
                head_h = max(int((py2 - py1) * HEAD_CROP_FRACTION), HEAD_CROP_MIN_PX)
                pad_x = max(4, int((px2 - px1) * 0.05))
                hx1 = max(0, px1 - pad_x)
                hx2 = min(w_img, px2 + pad_x)
                hy1 = max(0, py1)
                hy2 = min(h_img, py1 + head_h)

                # Determine rider color based on helmet status
                # (simplified — we can't re-run classification here,
                #  so we use the aggregate to decide styling)
                rider_color = CLR_NO_HELMET if without_h > 0 else CLR_RIDER
                label = "No Helmet" if without_h > 0 else "Helmet"

                cv2.rectangle(out, (px1, py1), (px2, py2), rider_color, 1)
                cv2.rectangle(out, (hx1, hy1), (hx2, hy2), rider_color, 1)
                put_label(out, label, px1, py1 - 5, (255, 255, 255), rider_color, scale=0.42, thickness=1)

        # Draw seatbelt violation boxes for cars
        if not is_bike and no_seatbelt > 0:
            # Draw the no-seatbelt boxes that fall inside this car
            no_seatbelt_boxes = annot.get("no_seatbelt_boxes", [])
            for sb in no_seatbelt_boxes:
                sb_cx = (sb[0] + sb[2]) / 2
                sb_cy = (sb[1] + sb[3]) / 2
                if x1 <= sb_cx <= x2 and y1 <= sb_cy <= y2:
                    cv2.rectangle(
                        out,
                        (int(sb[0]), int(sb[1])),
                        (int(sb[2]), int(sb[3])),
                        CLR_SEATBELT, 2,
                    )
                    put_label(
                        out, "No Seatbelt",
                        int(sb[0]), int(sb[1]) - 5,
                        (255, 255, 255), CLR_SEATBELT, scale=0.45,
                    )

        # Draw license plate box
        if plate_box_abs:
            px1, py1, px2, py2 = plate_box_abs
            cv2.rectangle(out, (px1, py1), (px2, py2), CLR_PLATE, 2)
            put_label(out, plate_text, px1, py2 + 16, (255, 255, 255), CLR_PLATE, scale=0.48, thickness=1)

    # ── Summary banner ───────────────────────────────────────────────────────
    total_all = total_violations

    banner_h = 38
    banner = np.zeros((banner_h, w_img, 3), dtype=np.uint8)
    if total_all > 0:
        banner[:] = (0, 0, 160)
        parts = [f"  {total_violations} VIOLATION(S)"]
        parts.append(f"Bikes: {total_bikes}")
        parts.append(f"Cars: {total_cars}")
        summary = "  |  ".join(parts)
    else:
        banner[:] = (0, 120, 0)
        summary = f"  NO VIOLATIONS  |  Bikes: {total_bikes}  |  Cars: {total_cars}"

    cv2.putText(banner, summary, (10, 26), FONT, 0.65, (255, 255, 255), 1, cv2.LINE_AA)
    out = np.vstack([banner, out])

    return out