utils.py

"""
Utility Functions

Kumpulan fungsi helper untuk drawing, resize, dan kalkulasi FPS.
Dipakai oleh app.py untuk rendering hasil deteksi dan tracking ke frame video.
"""

import cv2
import numpy as np
import time


# palette warna yang cukup kontras satu sama lain (BGR format)
COLOR_PALETTE = [
    (255, 150, 50),     # biru muda
    (50, 255, 50),      # hijau
    (50, 100, 255),     # oranye/merah
    (255, 50, 200),     # ungu/pink
    (0, 255, 255),      # kuning
    (255, 255, 0),      # cyan
    (128, 0, 255),      # magenta
    (0, 165, 255),      # oranye
]

# warna default kalau kelas tidak dikenali
DEFAULT_COLOR = (200, 200, 200)

# cache warna per class name supaya konsisten
_color_cache = {}


def get_class_color(class_name):
    """Ambil warna untuk class tertentu, konsisten selama runtime."""
    if class_name not in _color_cache:
        idx = len(_color_cache) % len(COLOR_PALETTE)
        _color_cache[class_name] = COLOR_PALETTE[idx]
    return _color_cache[class_name]


def draw_detections(frame, detections):
    """
    Gambar bounding box dan label pada frame.
    
    Args:
        frame: numpy array (BGR image)
        detections: list of dict dari detector.detect()
            setiap dict punya: bbox, confidence, class_name
    
    Returns:
        frame yang sudah di-annotate
    """
    annotated = frame.copy()

    for det in detections:
        bbox = det["bbox"]
        x1, y1, x2, y2 = [int(v) for v in bbox]
        conf = det["confidence"]
        cls_name = det["class_name"]

        color = get_class_color(cls_name)

        # gambar rectangle
        cv2.rectangle(annotated, (x1, y1), (x2, y2), color, 2)

        # buat label
        label = f"{cls_name} {conf:.2f}"
        label_size, _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 1)

        # background label
        cv2.rectangle(
            annotated,
            (x1, y1 - label_size[1] - 6),
            (x1 + label_size[0] + 4, y1),
            color,
            -1
        )

        # text label
        cv2.putText(
            annotated, label,
            (x1 + 2, y1 - 4),
            cv2.FONT_HERSHEY_SIMPLEX,
            0.5, (255, 255, 255), 1
        )

    return annotated


def draw_tracking(frame, tracked_objects):
    """
    Gambar bounding box dengan track ID pada frame.
    
    Args:
        frame: numpy array (BGR)
        tracked_objects: list of dict dari tracker.update()
            setiap dict punya: track_id, bbox, class_name, confidence
    
    Returns:
        frame yang sudah di-annotate
    """
    annotated = frame.copy()

    for obj in tracked_objects:
        bbox = obj["bbox"]
        x1, y1, x2, y2 = [int(v) for v in bbox]
        track_id = obj["track_id"]
        cls_name = obj["class_name"]
        conf = obj["confidence"]

        color = get_class_color(cls_name)

        # gambar rectangle
        cv2.rectangle(annotated, (x1, y1), (x2, y2), color, 2)

        # label dengan ID
        label = f"ID:{track_id} {cls_name} {conf:.2f}"
        label_size, _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 1)

        cv2.rectangle(
            annotated,
            (x1, y1 - label_size[1] - 6),
            (x1 + label_size[0] + 4, y1),
            color,
            -1
        )

        cv2.putText(
            annotated, label,
            (x1 + 2, y1 - 4),
            cv2.FONT_HERSHEY_SIMPLEX,
            0.5, (255, 255, 255), 1
        )

        # gambar center point
        cx, cy = obj["center"]
        cv2.circle(annotated, (int(cx), int(cy)), 4, color, -1)

        # gambar trajectory (beberapa titik terakhir)
        history = obj.get("center_history", [])
        if len(history) > 1:
            for i in range(1, len(history)):
                pt1 = (int(history[i-1][0]), int(history[i-1][1]))
                pt2 = (int(history[i][0]), int(history[i][1]))
                # fade effect: makin lama makin transparan
                thickness = max(1, int(2 * (i / len(history))))
                cv2.line(annotated, pt1, pt2, color, thickness)

    return annotated


def draw_counting_line(frame, line_position_ratio, count_text=""):
    """
    Gambar garis virtual horizontal pada frame.
    
    Args:
        frame: numpy array
        line_position_ratio: rasio posisi garis (0.0 - 1.0)
        count_text: text tambahan yang ditampilkan di dekat garis
    
    Returns:
        frame yang sudah digambar garisnya
    """
    annotated = frame.copy()
    h, w = frame.shape[:2]
    line_y = int(h * line_position_ratio)

    # garis utama (merah, tebal)
    cv2.line(annotated, (0, line_y), (w, line_y), (0, 0, 255), 2)

    # garis dashed effect (biar kelihatan lebih jelas)
    dash_length = 20
    for x in range(0, w, dash_length * 2):
        x_end = min(x + dash_length, w)
        cv2.line(annotated, (x, line_y), (x_end, line_y), (0, 255, 255), 3)

    # label "COUNTING LINE"
    cv2.putText(
        annotated, "COUNTING LINE",
        (10, line_y - 10),
        cv2.FONT_HERSHEY_SIMPLEX,
        0.6, (0, 255, 255), 2
    )

    if count_text:
        cv2.putText(
            annotated, count_text,
            (10, line_y + 25),
            cv2.FONT_HERSHEY_SIMPLEX,
            0.6, (0, 255, 255), 2
        )

    return annotated


def draw_polygon_region(frame, polygon_points):
    """
    Gambar polygon region pada frame.
    
    Args:
        frame: numpy array
        polygon_points: list of (x, y) tuples
    
    Returns:
        frame dengan overlay polygon
    """
    annotated = frame.copy()

    if not polygon_points or len(polygon_points) < 3:
        return annotated

    pts = np.array(polygon_points, dtype=np.int32)

    # gambar filled polygon semi-transparan
    overlay = annotated.copy()
    cv2.fillPoly(overlay, [pts], (0, 255, 0))
    annotated = cv2.addWeighted(overlay, 0.2, annotated, 0.8, 0)

    # gambar border polygon
    cv2.polylines(annotated, [pts], isClosed=True, color=(0, 255, 0), thickness=2)

    # label
    cx = int(np.mean([p[0] for p in polygon_points]))
    cy = int(np.mean([p[1] for p in polygon_points]))
    cv2.putText(
        annotated, "COUNTING REGION",
        (cx - 80, cy),
        cv2.FONT_HERSHEY_SIMPLEX,
        0.6, (0, 255, 0), 2
    )

    return annotated


def draw_stats_overlay(frame, stats):
    """
    Gambar overlay statistik di pojok kiri atas frame.
    
    Args:
        frame: numpy array
        stats: dict berisi informasi yang mau ditampilkan
            contoh: {"FPS": "24.5", "Total": "15", "Car": "8", ...}
    
    Returns:
        frame dengan overlay stats
    """
    annotated = frame.copy()
    h, w = frame.shape[:2]

    # background semi-transparan
    overlay = annotated.copy()
    box_h = 30 + len(stats) * 25
    cv2.rectangle(overlay, (5, 5), (200, box_h), (0, 0, 0), -1)
    annotated = cv2.addWeighted(overlay, 0.6, annotated, 0.4, 0)

    # render setiap stat
    y_offset = 25
    for key, value in stats.items():
        text = f"{key}: {value}"
        cv2.putText(
            annotated, text,
            (15, y_offset),
            cv2.FONT_HERSHEY_SIMPLEX,
            0.5, (255, 255, 255), 1
        )
        y_offset += 25

    return annotated


def resize_frame(frame, max_width=1280):
    """
    Resize frame supaya tidak terlalu besar untuk display.
    Menjaga aspect ratio.
    
    Args:
        frame: numpy array
        max_width: lebar maksimum
    
    Returns:
        frame yang sudah di-resize (atau frame asli kalau sudah cukup kecil)
    """
    h, w = frame.shape[:2]

    if w <= max_width:
        return frame

    scale = max_width / w
    new_w = int(w * scale)
    new_h = int(h * scale)

    resized = cv2.resize(frame, (new_w, new_h), interpolation=cv2.INTER_AREA)
    return resized


def calculate_fps(start_time, frame_count):
    """
    Hitung FPS berdasarkan waktu mulai dan jumlah frame.
    
    Args:
        start_time: waktu mulai (dari time.time())
        frame_count: jumlah frame yang sudah diproses
    
    Returns:
        float: FPS value
    """
    elapsed = time.time() - start_time
    if elapsed <= 0 or frame_count <= 0:
        return 0.0
    return frame_count / elapsed


def format_time(seconds):
    """
    Format detik ke string mm:ss.
    
    Args:
        seconds: float, durasi dalam detik
    
    Returns:
        str: formatted time string
    """
    minutes = int(seconds) // 60
    secs = int(seconds) % 60
    return f"{minutes:02d}:{secs:02d}"