app.py

"""Korview AI — Perimeter Security Demo

Upload a video or use the sample to see YOLO-based intrusion detection
with configurable security zones.
"""

from __future__ import annotations

import json
import tempfile
from pathlib import Path

import cv2
import gradio as gr
import numpy as np
import torch
from ultralytics import YOLO

# Load model once at startup, use GPU if available
DEVICE = "cuda:0" if torch.cuda.is_available() else "cpu"
MODEL = YOLO("yolov8n.pt")
print(f"YOLO loaded on {DEVICE} (CUDA available: {torch.cuda.is_available()})")

COCO_NAMES: dict[int, str] = {
    0: "person", 1: "bicycle", 2: "car", 3: "motorcycle",
    5: "bus", 7: "truck", 16: "dog", 17: "horse",
}

THREAT_COLORS = {
    "CRITICAL": (0, 0, 255),
    "HIGH": (0, 100, 255),
    "MEDIUM": (0, 200, 255),
    "LOW": (200, 200, 0),
}


def point_in_polygon(px: float, py: float, polygon: list[list[float]]) -> bool:
    n = len(polygon)
    inside = False
    j = n - 1
    for i in range(n):
        xi, yi = polygon[i]
        xj, yj = polygon[j]
        if ((yi > py) != (yj > py)) and (px < (xj - xi) * (py - yi) / (yj - yi) + xi):
            inside = not inside
        j = i
    return inside


def parse_zones(zones_json: str) -> list[dict]:
    try:
        zones = json.loads(zones_json)
        if isinstance(zones, dict) and "zones" in zones:
            zones = zones["zones"]
        return zones
    except (json.JSONDecodeError, TypeError):
        return []


def draw_zones(frame: np.ndarray, zones: list[dict]) -> np.ndarray:
    h, w = frame.shape[:2]
    overlay = frame.copy()
    for zone in zones:
        polygon = zone.get("polygon", [])
        threat = zone.get("threat_level", "MEDIUM")
        name = zone.get("name", zone.get("zone_id", "Zone"))
        color = THREAT_COLORS.get(threat, (0, 200, 255))

        pts = np.array([[int(p[0] * w), int(p[1] * h)] for p in polygon], np.int32)
        cv2.fillPoly(overlay, [pts], color)
        cv2.polylines(frame, [pts], True, color, 2)
        cv2.putText(frame, f"{name} [{threat}]",
                    (pts[0][0], pts[0][1] - 8),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)

    cv2.addWeighted(overlay, 0.15, frame, 0.85, 0, frame)
    return frame


def check_zone(cx: float, cy: float, zones: list[dict]) -> tuple[str | None, str]:
    best_zone = None
    best_threat = "INFO"
    threat_order = {"INFO": 0, "LOW": 1, "MEDIUM": 2, "HIGH": 3, "CRITICAL": 4}
    for zone in zones:
        polygon = zone.get("polygon", [])
        threat = zone.get("threat_level", "MEDIUM")
        if point_in_polygon(cx, cy, polygon):
            if threat_order.get(threat, 0) > threat_order.get(best_threat, 0):
                best_zone = zone.get("name", zone.get("zone_id"))
                best_threat = threat
    return best_zone, best_threat


def process_video(
    video_path: str,
    confidence: float,
    zones_json: str,
    max_seconds: float,
    process_every_n: int,
) -> tuple[str, str]:
    if not video_path:
        return None, "No video provided"

    zones = parse_zones(zones_json)

    cap = cv2.VideoCapture(video_path)
    if not cap.isOpened():
        return None, "Failed to open video"

    fps = cap.get(cv2.CAP_PROP_FPS) or 30
    w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    max_frames = min(int(max_seconds * fps), total_frames)

    # Resize if large
    scale = 1.0
    if w > 640:
        scale = 640 / w
        w_out, h_out = 640, int(h * scale)
    else:
        w_out, h_out = w, h

    out_path = tempfile.mktemp(suffix=".mp4")
    fourcc = cv2.VideoWriter_fourcc(*"mp4v")
    out_fps = fps / process_every_n  # output at reduced fps
    writer = cv2.VideoWriter(out_path, fourcc, out_fps, (w_out, h_out))

    events_log = []
    frame_count = 0
    processed = 0
    classes = [0, 1, 2, 3, 5, 7, 16, 17]
    last_detections = []  # reuse for skipped frames

    try:
        while cap.isOpened() and frame_count < max_frames:
            ret, frame = cap.read()
            if not ret:
                break

            frame_count += 1

            # Skip frames for speed
            if frame_count % process_every_n != 0:
                continue

            processed += 1

            # Resize
            if scale < 1.0:
                frame = cv2.resize(frame, (w_out, h_out))

            # Draw zones
            if zones:
                frame = draw_zones(frame, zones)

            # YOLO detection on GPU
            results = MODEL.track(
                frame, persist=True, conf=confidence,
                classes=classes, device=DEVICE, verbose=False,
            )

            for result in results:
                if result.boxes is None:
                    continue
                for box in result.boxes:
                    x1, y1, x2, y2 = map(int, box.xyxy[0].tolist())
                    cls_id = int(box.cls[0])
                    conf = float(box.conf[0])
                    track_id = int(box.id[0]) if box.id is not None else None
                    obj_name = COCO_NAMES.get(cls_id, f"class_{cls_id}")

                    # Normalized center
                    cx = ((x1 + x2) / 2) / w_out
                    cy = ((y1 + y2) / 2) / h_out

                    # Zone check
                    zone_name, threat = check_zone(cx, cy, zones) if zones else (None, "INFO")

                    # Draw bbox
                    color = THREAT_COLORS.get(threat, (0, 255, 0))
                    label = f"{obj_name} {conf:.0%}"
                    if track_id is not None:
                        label += f" #{track_id}"
                    if zone_name:
                        label += f" @{zone_name}"

                    cv2.rectangle(frame, (x1, y1), (x2, y2), color, 2)
                    (tw, th), _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.45, 1)
                    cv2.rectangle(frame, (x1, y1 - th - 6), (x1 + tw, y1), color, -1)
                    cv2.putText(frame, label, (x1, y1 - 4),
                                cv2.FONT_HERSHEY_SIMPLEX, 0.45, (255, 255, 255), 1)

                    # Log event
                    ts = frame_count / fps
                    if threat != "INFO":
                        events_log.append(
                            f"[{ts:6.1f}s] **{threat}** — {obj_name} (conf: {conf:.0%})"
                            f"{' in ' + zone_name if zone_name else ''}"
                            f"{' #' + str(track_id) if track_id else ''}"
                        )

            writer.write(frame)
    finally:
        cap.release()
        writer.release()

    # Summary
    summary = f"## Results\n\n"
    summary += f"Processed **{processed}** of {frame_count} frames "
    summary += f"({frame_count/fps:.1f}s video, every {process_every_n}{'st' if process_every_n == 1 else 'rd'} frame)\n\n"
    if events_log:
        summary += f"**{len(events_log)} security events detected:**\n\n"
        for line in events_log[-50:]:
            summary += f"- {line}\n"
        if len(events_log) > 50:
            summary += f"\n... and {len(events_log) - 50} more\n"
    else:
        summary += "No security events detected in configured zones.\n\n"
        summary += "*Tip: adjust zones to cover the area where objects appear, or lower the confidence threshold.*\n"

    return out_path, summary


DEFAULT_ZONES = json.dumps([
    {
        "zone_id": "fence_area",
        "name": "Fence Perimeter",
        "threat_level": "HIGH",
        "polygon": [[0.0, 0.0], [1.0, 0.0], [1.0, 0.5], [0.0, 0.5]]
    },
    {
        "zone_id": "restricted",
        "name": "Restricted Area",
        "threat_level": "CRITICAL",
        "polygon": [[0.2, 0.5], [0.8, 0.5], [0.8, 0.95], [0.2, 0.95]]
    }
], indent=2)


with gr.Blocks(
    title="Korview AI — Perimeter Security",
    theme=gr.themes.Base(primary_hue="red", neutral_hue="slate"),
) as demo:
    gr.Markdown(
        "# Korview AI — Perimeter Security Detection\n"
        "Upload a security camera video to detect intrusions with YOLO. "
        "Configure security zones to classify threats."
    )

    with gr.Row():
        with gr.Column(scale=1):
            video_input = gr.Video(label="Upload Security Camera Video", sources=["upload"])
            confidence = gr.Slider(
                minimum=0.1, maximum=0.9, value=0.4, step=0.05,
                label="Detection Confidence",
            )
            max_seconds = gr.Slider(
                minimum=5, maximum=120, value=30, step=5,
                label="Max Video Duration (seconds)",
                info="T4 GPU: ~300 fps. 30s video ≈ 3s processing.",
            )
            process_every_n = gr.Slider(
                minimum=1, maximum=10, value=1, step=1,
                label="Process Every Nth Frame",
                info="1 = every frame (GPU). Increase on CPU for speed.",
            )
            zones_input = gr.Code(
                value=DEFAULT_ZONES,
                language="json",
                label="Security Zones (JSON — normalized 0.0-1.0 coords)",
                lines=12,
            )
            run_btn = gr.Button("Run Detection", variant="primary", size="lg")

        with gr.Column(scale=1):
            video_output = gr.Video(label="Detection Results")
            events_output = gr.Markdown(label="Security Events")

    run_btn.click(
        fn=process_video,
        inputs=[video_input, confidence, zones_input, max_seconds, process_every_n],
        outputs=[video_output, events_output],
    )

demo.launch(ssr_mode=False)