video_processor.py

"""video_processor.py
Core YOLOv8 detection + DeepSORT multi-object tracking pipeline for whoMoved.
"""
from __future__ import annotations

import functools
import os
import shutil
import subprocess
import threading
import time
from typing import Callable

import cv2
import numpy as np
from deep_sort_realtime.deepsort_tracker import DeepSort
from ultralytics import YOLO

# ── Tuneable constants ────────────────────────────────────────────────────────
CONF_THRESHOLD: float = 0.40      # slightly higher threshold = fewer false positives to track
MAX_TRACK_AGE: int = 20           # frames a lost track is kept alive before deletion
N_INIT: int = 2                   # detections required before a track is confirmed (faster confirmation)
MAX_DET: int = 30                 # cap detections per frame for speed
INFER_SIZE: int = 320             # YOLO input size — 320 is ~4× faster than 640 on CPU with minor accuracy loss
MODEL_NAME: str = "yolov8n.pt"    # YOLOv8 Nano — fastest; swap for yolov8s/m for accuracy
PROGRESS_STEP: int = 5            # update progress bar every N frames

# ── COCO class names (all 80) ─────────────────────────────────────────────────
COCO_CLASSES: list[str] = [
    "person", "bicycle", "car", "motorcycle", "airplane", "bus", "train",
    "truck", "boat", "traffic light", "fire hydrant", "stop sign",
    "parking meter", "bench", "bird", "cat", "dog", "horse", "sheep", "cow",
    "elephant", "bear", "zebra", "giraffe", "backpack", "umbrella", "handbag",
    "tie", "suitcase", "frisbee", "skis", "snowboard", "sports ball", "kite",
    "baseball bat", "baseball glove", "skateboard", "surfboard",
    "tennis racket", "bottle", "wine glass", "cup", "fork", "knife", "spoon",
    "bowl", "banana", "apple", "sandwich", "orange", "broccoli", "carrot",
    "hot dog", "pizza", "donut", "cake", "chair", "couch", "potted plant",
    "bed", "dining table", "toilet", "tv", "laptop", "mouse", "remote",
    "keyboard", "cell phone", "microwave", "oven", "toaster", "sink",
    "refrigerator", "book", "clock", "vase", "scissors", "teddy bear",
    "hair drier", "toothbrush",
]


# ── Colour helpers ────────────────────────────────────────────────────────────

@functools.lru_cache(maxsize=256)
def _track_color(track_id: int) -> tuple[int, int, int]:
    """Return a deterministic, vibrant BGR colour for a given track ID (cached)."""
    hue = int((track_id * 47) % 180)
    hsv = np.array([[[hue, 240, 220]]], dtype=np.uint8)
    bgr = cv2.cvtColor(hsv, cv2.COLOR_HSV2BGR)[0][0]
    return (int(bgr[0]), int(bgr[1]), int(bgr[2]))


# ── Drawing helpers ───────────────────────────────────────────────────────────

def _draw_box(
    frame: np.ndarray,
    x1: int, y1: int, x2: int, y2: int,
    label: str,
    color: tuple[int, int, int],
) -> None:
    """Draw an anti-aliased bounding box with a filled label pill."""
    # Bounding box
    cv2.rectangle(frame, (x1, y1), (x2, y2), color, 2, cv2.LINE_AA)

    # Label background
    font, scale, thick = cv2.FONT_HERSHEY_SIMPLEX, 0.55, 1
    (tw, th), baseline = cv2.getTextSize(label, font, scale, thick)
    bg_y1 = max(y1 - th - baseline - 8, 0)
    cv2.rectangle(frame, (x1, bg_y1), (x1 + tw + 6, y1), color, cv2.FILLED)

    # Pick text colour that contrasts against the background
    luma = 0.114 * color[0] + 0.587 * color[1] + 0.299 * color[2]
    text_color = (10, 10, 10) if luma > 128 else (240, 240, 240)
    cv2.putText(
        frame, label,
        (x1 + 3, y1 - baseline - 3),
        font, scale, text_color, thick, cv2.LINE_AA,
    )


def _draw_hud(frame: np.ndarray, frame_idx: int, total: int, active: int) -> None:
    """Burn a small HUD line into the bottom-left of the frame."""
    h = frame.shape[0]
    text = f"Frame {frame_idx + 1}/{total}   Active IDs: {active}"
    cv2.putText(
        frame, text, (8, h - 10),
        cv2.FONT_HERSHEY_SIMPLEX, 0.5, (200, 200, 200), 1, cv2.LINE_AA,
    )


def _draw_zone_line(
    frame: np.ndarray,
    zone_y: int,
    crossings_in: int,
    crossings_out: int,
) -> None:
    """Draw a cyan tripwire line with crossing counts burned in."""
    w = frame.shape[1]
    cv2.line(frame, (0, zone_y), (w, zone_y), (0, 220, 220), 2, cv2.LINE_AA)
    label = f"Zone  \u2193{crossings_in}  \u2191{crossings_out}"
    cv2.putText(
        frame, label, (8, max(zone_y - 6, 14)),
        cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 220, 220), 1, cv2.LINE_AA,
    )


# ── FFmpeg re-encoding ────────────────────────────────────────────────────────

def _reencode_h264(raw: str, out: str) -> None:
    """
    Re-encode *raw* (mp4v) to H.264/MP4 so any browser can play it inline.
    Falls back to a plain file move when ffmpeg is unavailable.
    """
    if not shutil.which("ffmpeg"):
        shutil.move(raw, out)
        return
    try:
        subprocess.run(
            [
                "ffmpeg", "-y", "-i", raw,
                "-vcodec", "libx264", "-crf", "28",
                "-preset", "fast", "-pix_fmt", "yuv420p",
                out,
            ],
            check=True,
            capture_output=True,
        )
        os.remove(raw)
    except subprocess.CalledProcessError:
        # ffmpeg failed — use the raw file as-is
        shutil.move(raw, out)


# ── Public API ────────────────────────────────────────────────────────────────

def process_video(
    input_path: str,
    output_path: str,
    progress_callback: Callable[[float], None] | None = None,
    conf: float = CONF_THRESHOLD,
    class_filter: set | None = None,
    zone_y_frac: float | None = None,
    model_name: str = MODEL_NAME,
) -> dict:
    """
    Run YOLOv8 detection + DeepSORT tracking on every frame of a video.

    Parameters
    ----------
    input_path        : absolute path to the source video
    output_path       : absolute path for the annotated output MP4
    progress_callback : optional callable that receives a float in [0, 1]
    conf              : minimum detection confidence (0.0–1.0)
    class_filter      : set of class names to keep; None = all classes
    zone_y_frac       : horizontal tripwire at y = H * zone_y_frac; None = disabled
    model_name        : YOLOv8 model weights file

    Returns
    -------
    dict with keys
        total_objects       – unique track IDs seen
        total_frames        – frames processed
        duration            – video duration in seconds
        fps                 – source FPS
        track_details       – list[dict], one entry per track ID
        zone_crossings_in   – above→below crossings (None if zone disabled)
        zone_crossings_out  – below→above crossings (None if zone disabled)
    """
    # ── Load model & tracker ─────────────────────────────────────────────────
    model = YOLO(model_name)
    tracker = DeepSort(max_age=MAX_TRACK_AGE, n_init=N_INIT)

    # ── Open input video ─────────────────────────────────────────────────────
    cap = cv2.VideoCapture(input_path)
    if not cap.isOpened():
        raise ValueError(f"Cannot open video file: {input_path}")

    total_frames = max(int(cap.get(cv2.CAP_PROP_FRAME_COUNT)), 1)
    fps = cap.get(cv2.CAP_PROP_FPS) or 25.0
    W = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    H = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))

    # ── Open output writer (mp4v; will be re-encoded later) ──────────────────
    raw_path = output_path + ".raw.mp4"
    fourcc = cv2.VideoWriter_fourcc(*"mp4v")
    writer = cv2.VideoWriter(raw_path, fourcc, fps, (W, H))

    # Per-track accumulators
    track_stats: dict[int, dict] = {}
    class_cache: dict[int, str] = {}
    frame_idx = 0

    # Zone crossing accumulators
    last_side: dict[int, int] = {}
    crossings_in: int = 0
    crossings_out: int = 0

    # ── Frame loop ───────────────────────────────────────────────────────────
    while True:
        ret, frame = cap.read()
        if not ret:
            break

        # ── YOLOv8 detection ─────────────────────────────────────────────────
        results = model(frame, imgsz=INFER_SIZE, max_det=MAX_DET, verbose=False)[0]
        raw_dets: list[tuple] = []
        for box in results.boxes:
            box_conf = float(box.conf[0])
            if box_conf < conf:
                continue
            x1, y1, x2, y2 = box.xyxy[0].tolist()
            cls_name = model.names[int(box.cls[0])]
            if class_filter is not None and cls_name not in class_filter:
                continue
            raw_dets.append(([x1, y1, x2 - x1, y2 - y1], box_conf, cls_name))

        # ── DeepSORT tracking ─────────────────────────────────────────────────
        active_tracks = tracker.update_tracks(raw_dets, frame=frame)
        active_count = 0
        zone_y_px = int(H * zone_y_frac) if zone_y_frac is not None else None

        for track in active_tracks:
            if not track.is_confirmed():
                continue
            active_count += 1
            tid = int(track.track_id)

            ltrb = track.to_ltrb()
            x1 = max(0, int(ltrb[0]))
            y1 = max(0, int(ltrb[1]))
            x2 = min(W, int(ltrb[2]))
            y2 = min(H, int(ltrb[3]))

            det_cls = track.det_class
            if det_cls:
                class_cache[tid] = det_cls
            cls_name = class_cache.get(tid, "object")

            if tid not in track_stats:
                track_stats[tid] = {
                    "class": cls_name,
                    "first_frame": frame_idx,
                    "last_frame": frame_idx,
                    "frame_count": 1,
                }
            else:
                track_stats[tid]["last_frame"] = frame_idx
                track_stats[tid]["frame_count"] += 1

            # Zone crossing detection
            if zone_y_px is not None:
                cy = (y1 + y2) // 2
                side = 1 if cy >= zone_y_px else -1
                if tid in last_side and last_side[tid] != side:
                    if side == 1:
                        crossings_in += 1
                    else:
                        crossings_out += 1
                last_side[tid] = side

            color = _track_color(tid)
            _draw_box(frame, x1, y1, x2, y2, f"ID {tid}  {cls_name}", color)

        if zone_y_px is not None:
            _draw_zone_line(frame, zone_y_px, crossings_in, crossings_out)
        _draw_hud(frame, frame_idx, total_frames, active_count)
        writer.write(frame)
        frame_idx += 1

        if progress_callback and frame_idx % PROGRESS_STEP == 0:
            progress_callback(min(frame_idx / total_frames, 0.99))

    cap.release()
    writer.release()

    if progress_callback:
        progress_callback(1.0)

    # ── Re-encode for browser playback ───────────────────────────────────────
    _reencode_h264(raw_path, output_path)

    # ── Build and return statistics ──────────────────────────────────────────
    duration = frame_idx / fps
    track_details = sorted(
        [
            {
                "Track ID": tid,
                "Class": info["class"],
                "Tracking Duration (s)": round(info["frame_count"] / fps, 2),
                "First Frame": info["first_frame"] + 1,
                "Last Frame": info["last_frame"] + 1,
                "Frame Count": info["frame_count"],
            }
            for tid, info in track_stats.items()
        ],
        key=lambda r: r["Track ID"],
    )

    return {
        "total_objects": len(track_stats),
        "total_frames": frame_idx,
        "duration": round(duration, 2),
        "fps": round(fps, 2),
        "track_details": track_details,
        "zone_crossings_in": crossings_in if zone_y_frac is not None else None,
        "zone_crossings_out": crossings_out if zone_y_frac is not None else None,
    }


# ── Real-time processor ───────────────────────────────────────────────────────

class RealtimeProcessor:
    """
    Stateful, thread-safe YOLOv8 + DeepSORT processor for live webcam use.

    backend options
    ---------------
    "pytorch"      – plain PyTorch CPU inference (always works)
    "openvino_cpu" – OpenVINO on Intel CPU (faster graph fusion)
    "openvino_gpu" – OpenVINO on Intel iGPU (Iris / Iris Xe / Arc)
    """

    def __init__(self, backend: str = "pytorch", model_name: str = MODEL_NAME) -> None:
        _ov_dir = model_name.replace(".pt", "_openvino_model")
        if backend.startswith("openvino"):
            if not os.path.exists(_ov_dir):
                YOLO(model_name).export(format="openvino")
            self.model = YOLO(_ov_dir, task="detect")
            self._ov_device: str | None = "GPU" if backend == "openvino_gpu" else "CPU"
            # OpenVINO models have a static input shape baked in at export time (640).
            # Passing a different imgsz at runtime causes a shape-mismatch RuntimeError,
            # so we leave imgsz unset and let the model use its own compiled size.
            self._infer_size: int | None = None
        else:
            self.model = YOLO(model_name)
            self._ov_device = None
            self._infer_size = INFER_SIZE

        self.tracker = DeepSort(max_age=MAX_TRACK_AGE, n_init=N_INIT)
        self._class_cache: dict[int, str] = {}
        self._track_stats: dict[int, dict] = {}
        self.frame_count: int = 0
        self._lock = threading.Lock()

        # Mutable runtime config — updated live via update_config()
        self.conf_threshold: float = CONF_THRESHOLD
        self.allowed_classes: set | None = None   # None = all classes
        self.zone_y_frac: float | None = None     # None = no zone line

        # FPS tracking (exponential moving average)
        self._fps: float = 0.0
        self._t_last: float | None = None

        # Zone crossing state
        self._last_side: dict[int, int] = {}
        self._crossings_in: int = 0
        self._crossings_out: int = 0

    # ── public API ────────────────────────────────────────────────────────────

    def update_config(
        self,
        conf: float,
        classes: set | None,
        zone_y_frac: float | None,
    ) -> None:
        """Thread-safe update of runtime detection / zone config."""
        with self._lock:
            self.conf_threshold = conf
            self.allowed_classes = classes
            self.zone_y_frac = zone_y_frac

    def process_frame(self, frame: np.ndarray) -> tuple[np.ndarray, dict]:
        """Detect, track, and annotate a single BGR frame. Returns (annotated, stats)."""
        t0 = time.perf_counter()
        h, w = frame.shape[:2]

        # Read config — CPython GIL makes simple attr reads effectively atomic
        conf_threshold = self.conf_threshold
        allowed_classes = self.allowed_classes
        zone_y_frac = self.zone_y_frac

        # ── YOLOv8 detection ──────────────────────────────────────────────────
        kwargs: dict = {"verbose": False, "max_det": MAX_DET}
        if self._infer_size is not None:
            kwargs["imgsz"] = self._infer_size
        if self._ov_device:
            kwargs["device"] = self._ov_device
        results = self.model(frame, **kwargs)[0]

        raw_dets: list[tuple] = []
        for box in results.boxes:
            box_conf = float(box.conf[0])
            if box_conf < conf_threshold:
                continue
            x1, y1, x2, y2 = box.xyxy[0].tolist()
            cls_name = self.model.names[int(box.cls[0])]
            if allowed_classes is not None and cls_name not in allowed_classes:
                continue
            raw_dets.append(([x1, y1, x2 - x1, y2 - y1], box_conf, cls_name))

        # ── DeepSORT tracking (lock protects mutable state) ───────────────────
        with self._lock:
            active_tracks = self.tracker.update_tracks(raw_dets, frame=frame)
            active_count = 0
            zone_y_px = int(h * zone_y_frac) if zone_y_frac is not None else None

            for track in active_tracks:
                if not track.is_confirmed():
                    continue
                active_count += 1
                tid = int(track.track_id)

                ltrb = track.to_ltrb()
                x1 = max(0, int(ltrb[0]))
                y1 = max(0, int(ltrb[1]))
                x2 = min(w, int(ltrb[2]))
                y2 = min(h, int(ltrb[3]))

                if track.det_class:
                    self._class_cache[tid] = track.det_class
                cls_name = self._class_cache.get(tid, "object")

                if tid not in self._track_stats:
                    self._track_stats[tid] = {"class": cls_name, "frame_count": 1}
                else:
                    self._track_stats[tid]["frame_count"] += 1

                # Zone crossing detection
                if zone_y_px is not None:
                    cy = (y1 + y2) // 2
                    side = 1 if cy >= zone_y_px else -1
                    if tid in self._last_side and self._last_side[tid] != side:
                        if side == 1:
                            self._crossings_in += 1
                        else:
                            self._crossings_out += 1
                    self._last_side[tid] = side

                _draw_box(frame, x1, y1, x2, y2, f"ID {tid}  {cls_name}", _track_color(tid))

            # Zone line overlay
            if zone_y_px is not None:
                _draw_zone_line(frame, zone_y_px, self._crossings_in, self._crossings_out)

            # FPS — exponential moving average over inference calls
            if self._t_last is not None:
                dt = t0 - self._t_last
                if dt > 0:
                    inst = 1.0 / dt
                    self._fps = self._fps * 0.85 + inst * 0.15 if self._fps > 0 else inst
            self._t_last = t0

            # HUD
            cv2.putText(
                frame,
                f"FPS {self._fps:.1f}  |  Active: {active_count}  |  Total IDs: {len(self._track_stats)}",
                (8, h - 10),
                cv2.FONT_HERSHEY_SIMPLEX, 0.5, (200, 200, 200), 1, cv2.LINE_AA,
            )
            self.frame_count += 1

            stats = {
                "active_count": active_count,
                "total_ids": len(self._track_stats),
                "frame_count": self.frame_count,
                "fps": round(self._fps, 1),
                "crossings_in": self._crossings_in,
                "crossings_out": self._crossings_out,
            }

        return frame, stats

    def reset(self) -> None:
        with self._lock:
            self.tracker = DeepSort(max_age=MAX_TRACK_AGE, n_init=N_INIT)
            self._class_cache.clear()
            self._track_stats.clear()
            self.frame_count = 0
            self._fps = 0.0
            self._t_last = None
            self._last_side.clear()
            self._crossings_in = 0
            self._crossings_out = 0

    def get_summary(self) -> dict:
        with self._lock:
            return {
                "total_ids": len(self._track_stats),
                "frame_count": self.frame_count,
                "fps": round(self._fps, 1),
                "crossings_in": self._crossings_in,
                "crossings_out": self._crossings_out,
                "track_details": sorted(
                    [
                        {
                            "Track ID": tid,
                            "Class": v["class"],
                            "Frame Count": v["frame_count"],
                        }
                        for tid, v in self._track_stats.items()
                    ],
                    key=lambda r: r["Track ID"],
                ),
            }