app.py

"""
Fight / Violence Detector - MVP
Zero-shot video action detection using X-CLIP.
- Upload-video tab
- Live-webcam tab
- Telegram alerts + browser alarm sound on trigger
"""
import os
import time
import sqlite3
import datetime
import threading
from pathlib import Path

import numpy as np
import torch
import cv2
import gradio as gr
import requests
from dotenv import load_dotenv
from transformers import VideoMAEImageProcessor, VideoMAEForVideoClassification

load_dotenv()

MODEL_ID = os.environ.get(
    "MODEL_ID",
    "archit11/videomae-base-finetuned-fight-nofight-subset2",
)
ALERT_THRESHOLD = float(os.environ.get("ALERT_THRESHOLD", "0.60"))
NUM_FRAMES = 16
DB_PATH = Path("events.db")
ALERT_COOLDOWN_SEC = 10

TELEGRAM_TOKEN = os.environ.get("TELEGRAM_TOKEN")
TELEGRAM_CHAT_ID = os.environ.get("TELEGRAM_CHAT_ID")

device = "cuda" if torch.cuda.is_available() else "cpu"
processor = VideoMAEImageProcessor.from_pretrained(MODEL_ID)
model = VideoMAEForVideoClassification.from_pretrained(MODEL_ID).to(device).eval()

ID2LABEL = {int(k): v for k, v in model.config.id2label.items()}


def _find_alert_index():
    for idx, label in ID2LABEL.items():
        norm = label.lower().replace(" ", "").replace("_", "").replace("-", "")
        if norm in ("fight", "violence"):
            return idx
    return 0


ALERT_IDX = _find_alert_index()
ALERT_LABEL = ID2LABEL[ALERT_IDX]
print(f"[config] model={MODEL_ID}  labels={ID2LABEL}  alert_on={ALERT_LABEL}")

_last_alert_ts = 0.0
_alert_lock = threading.Lock()


def init_db():
    con = sqlite3.connect(DB_PATH)
    con.execute(
        "CREATE TABLE IF NOT EXISTS events "
        "(ts TEXT, label TEXT, confidence REAL, source TEXT)"
    )
    con.commit()
    con.close()


def log_event(label, confidence, source):
    con = sqlite3.connect(DB_PATH)
    con.execute(
        "INSERT INTO events VALUES (?,?,?,?)",
        (datetime.datetime.now().isoformat(timespec="seconds"), label, confidence, source),
    )
    con.commit()
    con.close()


def recent_events(limit=20):
    con = sqlite3.connect(DB_PATH)
    rows = con.execute(
        "SELECT ts, label, confidence, source FROM events ORDER BY ts DESC LIMIT ?",
        (limit,),
    ).fetchall()
    con.close()
    return rows


def send_telegram(message):
    if not TELEGRAM_TOKEN or not TELEGRAM_CHAT_ID:
        return
    try:
        requests.post(
            f"https://api.telegram.org/bot{TELEGRAM_TOKEN}/sendMessage",
            data={"chat_id": TELEGRAM_CHAT_ID, "text": message},
            timeout=5,
        )
    except requests.RequestException:
        pass


def generate_alarm():
    sr = 22050
    duration = 1.5
    t = np.linspace(0, duration, int(sr * duration), endpoint=False)
    wave = np.where(
        (t * 4) % 1 < 0.5,
        np.sin(2 * np.pi * 880 * t),
        np.sin(2 * np.pi * 660 * t),
    )
    return sr, (wave * 0.4 * 32767).astype(np.int16)


def should_alert(confidence):
    global _last_alert_ts
    if confidence < ALERT_THRESHOLD:
        return False
    with _alert_lock:
        now = time.time()
        if now - _last_alert_ts < ALERT_COOLDOWN_SEC:
            return False
        _last_alert_ts = now
    return True


def run_inference(frames):
    inputs = processor(list(frames), return_tensors="pt").to(device)
    with torch.no_grad():
        outputs = model(**inputs)
        probs = outputs.logits.softmax(dim=-1)[0].cpu().numpy()
    return {ID2LABEL[i]: float(p) for i, p in enumerate(probs)}


def trigger_alert(confidence, source):
    log_event(ALERT_LABEL, confidence, source)
    send_telegram(f"ALERT: fight detected ({confidence:.0%}) - source: {source}")
    return generate_alarm()


def sample_frames_from_video(video_path, num_frames=NUM_FRAMES):
    cap = cv2.VideoCapture(video_path)
    total = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) or 1
    indices = set(np.linspace(0, max(total - 1, 0), num_frames).astype(int).tolist())
    frames, i = [], 0
    while True:
        ok, frame = cap.read()
        if not ok:
            break
        if i in indices:
            frames.append(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
        i += 1
    cap.release()
    while len(frames) < num_frames and frames:
        frames.append(frames[-1])
    return np.stack(frames)


def detect_uploaded(video_path):
    if video_path is None:
        return {}, "Upload a video", None, recent_events()
    frames = sample_frames_from_video(video_path)
    scores = run_inference(frames)
    alert_conf = scores[ALERT_LABEL]
    alarm = None
    if should_alert(alert_conf):
        alarm = trigger_alert(alert_conf, source=os.path.basename(video_path))
        verdict = f"ALERT: fight detected ({alert_conf:.0%})"
    else:
        verdict = f"Normal scene ({alert_conf:.0%} fight confidence)"
    return scores, verdict, alarm, recent_events()


def stream_frame(frame, buffer):
    if frame is None:
        return buffer or [], gr.update(), "waiting for webcam...", None, gr.update()
    buffer = (buffer or []) + [frame]
    if len(buffer) < NUM_FRAMES:
        return (
            buffer,
            gr.update(),
            f"collecting frames {len(buffer)}/{NUM_FRAMES}",
            None,
            gr.update(),
        )
    scores = run_inference(np.stack(buffer))
    alert_conf = scores[ALERT_LABEL]
    alarm = None
    verdict = f"Normal ({alert_conf:.0%} fight confidence)"
    if should_alert(alert_conf):
        alarm = trigger_alert(alert_conf, source="webcam")
        verdict = f"ALERT: fight detected ({alert_conf:.0%})"
    return [], scores, verdict, alarm, recent_events()


def monitor_stream(url):
    """Generator that opens an RTSP/HTTP stream and yields inference results in a loop."""
    if not url or not url.strip():
        yield None, {}, "Enter a stream URL and click Start", None, recent_events()
        return

    cap = cv2.VideoCapture(url.strip())
    if not cap.isOpened():
        yield None, {}, f"Cannot open stream: {url}", None, recent_events()
        return

    buffer = []
    try:
        while True:
            ok, frame = cap.read()
            if not ok:
                yield None, gr.update(), "Stream ended or connection lost", None, gr.update()
                return
            rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            buffer.append(rgb)
            if len(buffer) < NUM_FRAMES:
                yield rgb, gr.update(), f"collecting frames {len(buffer)}/{NUM_FRAMES}", None, gr.update()
                continue
            scores = run_inference(np.stack(buffer))
            buffer = []
            alert_conf = scores[ALERT_LABEL]
            alarm = None
            verdict = f"Normal ({alert_conf:.0%} fight confidence)"
            if should_alert(alert_conf):
                alarm = trigger_alert(alert_conf, source=url[:80])
                verdict = f"ALERT: fight detected ({alert_conf:.0%})"
            yield rgb, scores, verdict, alarm, recent_events()
    finally:
        cap.release()


init_db()

tg_status = (
    "configured"
    if (TELEGRAM_TOKEN and TELEGRAM_CHAT_ID)
    else "NOT configured — set TELEGRAM_TOKEN and TELEGRAM_CHAT_ID in .env"
)

HIDE_FOOTER_CSS = """
footer { visibility: hidden !important; }
"""

with gr.Blocks(title="Fight Detector MVP", css=HIDE_FOOTER_CSS) as demo:
    gr.Markdown("# Fight / Violence Detector")
    gr.Markdown(f"**Telegram alerts:** {tg_status}  |  **Cooldown:** {ALERT_COOLDOWN_SEC}s between alerts")

    with gr.Tabs():
        with gr.Tab("Upload Video"):
            with gr.Row():
                with gr.Column():
                    video_in = gr.Video(label="Input video", sources=["upload"])
                    run_btn = gr.Button("Analyze", variant="primary")
                with gr.Column():
                    verdict1 = gr.Textbox(label="Result", interactive=False)
                    scores1 = gr.Label(label="Scores", num_top_classes=3)
                    alarm1 = gr.Audio(label="Alarm", autoplay=True)
            events_table1 = gr.Dataframe(
                headers=["Time", "Label", "Confidence", "Source"],
                label="Recent alerts",
                value=recent_events(),
                interactive=False,
            )
            run_btn.click(
                detect_uploaded,
                inputs=video_in,
                outputs=[scores1, verdict1, alarm1, events_table1],
            )

        with gr.Tab("Live Webcam"):
            gr.Markdown("Grant webcam permission, then just watch. Inference runs every 8 frames.")
            with gr.Row():
                with gr.Column():
                    cam = gr.Image(
                        sources=["webcam"],
                        streaming=True,
                        label="Live stream",
                        type="numpy",
                    )
                with gr.Column():
                    verdict2 = gr.Textbox(label="Result", interactive=False)
                    scores2 = gr.Label(label="Scores", num_top_classes=3)
                    alarm2 = gr.Audio(label="Alarm", autoplay=True)

            buffer_state = gr.State([])
            events_table2 = gr.Dataframe(
                headers=["Time", "Label", "Confidence", "Source"],
                label="Recent alerts",
                value=recent_events(),
                interactive=False,
            )
            cam.stream(
                stream_frame,
                inputs=[cam, buffer_state],
                outputs=[buffer_state, scores2, verdict2, alarm2, events_table2],
            )

        with gr.Tab("IP Camera / Mobile"):
            gr.Markdown(
                "Paste a stream URL from a security camera (RTSP) or phone.\n\n"
                "**Examples:**\n"
                "- Android `IP Webcam` app: `http://192.168.1.10:8080/video`\n"
                "- Security camera RTSP: `rtsp://user:pass@192.168.1.20:554/stream`\n"
                "- Public HLS/RTSP test stream\n"
            )
            with gr.Row():
                with gr.Column():
                    url_in = gr.Textbox(
                        label="Stream URL",
                        placeholder="rtsp://user:pass@192.168.1.20:554/stream",
                    )
                    with gr.Row():
                        start_btn = gr.Button("Start monitoring", variant="primary")
                        stop_btn = gr.Button("Stop")
                    preview = gr.Image(label="Live preview", type="numpy", streaming=False)
                with gr.Column():
                    verdict3 = gr.Textbox(label="Result", interactive=False)
                    scores3 = gr.Label(label="Scores", num_top_classes=3)
                    alarm3 = gr.Audio(label="Alarm", autoplay=True)

            events_table3 = gr.Dataframe(
                headers=["Time", "Label", "Confidence", "Source"],
                label="Recent alerts",
                value=recent_events(),
                interactive=False,
            )

            stream_event = start_btn.click(
                monitor_stream,
                inputs=[url_in],
                outputs=[preview, scores3, verdict3, alarm3, events_table3],
            )
            stop_btn.click(fn=None, inputs=None, outputs=None, cancels=[stream_event])


if __name__ == "__main__":
    host = "0.0.0.0" if os.environ.get("SPACE_ID") else "127.0.0.1"
    demo.queue().launch(server_name=host, server_port=7860, show_api=False)