app.py

import io
import time
 
import spaces
import cv2
import numpy as np
import pandas as pd
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
from PIL import Image
import torch
import gradio as gr
from huggingface_hub import hf_hub_download
 
from model import FallDetector, EFFICIENTNET_DIM
 
 
# Configuration
 
REPO_ID       = "beaunix/aegis-fall-detector"
CKPT_FILENAME = "fall_detector_best.pt"
 
N_FRAMES       = 16
IMG_SIZE       = 224
FALL_THRESHOLD = 0.65
MAX_DURATION   = 45.0   # seconds; longer videos are rejected
 
MEAN = np.array([0.485, 0.456, 0.406], dtype=np.float32)
STD  = np.array([0.229, 0.224, 0.225], dtype=np.float32)
 
# Cyberpunk HUD palette - orange edition
BG      = "#000000"
PANEL   = "#0a0a0a"
ORANGE  = "#FF8C00"
AMBER   = "#FFB347"
SILVER  = "#C0C0C0"
RED     = "#FF3030"
GREEN   = "#00FF9C"
GRID    = "#1c1c1c"
 
 
 
# Model loading (CPU, once at startup)
 
def load_model():
    ckpt_path = hf_hub_download(repo_id=REPO_ID, filename=CKPT_FILENAME)
    state = torch.load(ckpt_path, map_location="cpu", weights_only=False)
 
    net = FallDetector(pretrained_backbone=False)
    missing, unexpected = net.load_state_dict(state, strict=False)
 
    real_missing = [k for k in missing if not k.endswith("num_batches_tracked")]
    total_keys   = len(net.state_dict())
    loaded_ratio = (total_keys - len(real_missing)) / total_keys
 
    print(f"[LOAD] Loaded {loaded_ratio*100:.1f}% of params "
          f"({len(real_missing)} missing, {len(unexpected)} unexpected).")
 
    if loaded_ratio < 0.95:
        raise RuntimeError(
            "Checkpoint keys do not match the model. Weights were NOT loaded "
            f"correctly (only {loaded_ratio*100:.1f}% matched). "
            f"First missing: {real_missing[:5]} | First unexpected: {unexpected[:5]}"
        )
 
    net.eval()
    return net
 
 
print("[INIT] Loading Fall Detector model...")
MODEL = load_model()
print("[INIT] Model ready (CPU).")
 
 
# Preprocessing - letterbox (matches training ETL exactly)
 
def letterbox_frame(frame_bgr: np.ndarray, target: int = IMG_SIZE) -> np.ndarray:
    """Resize keeping aspect ratio, pad with black to target x target. Returns RGB."""
    h, w  = frame_bgr.shape[:2]
    scale = target / max(h, w)
    new_w = int(w * scale)
    new_h = int(h * scale)
 
    resized = cv2.resize(frame_bgr, (new_w, new_h), interpolation=cv2.INTER_LINEAR)
 
    canvas = np.zeros((target, target, 3), dtype=np.uint8)
    pad_top  = (target - new_h) // 2
    pad_left = (target - new_w) // 2
    canvas[pad_top:pad_top + new_h, pad_left:pad_left + new_w] = resized
 
    return cv2.cvtColor(canvas, cv2.COLOR_BGR2RGB)
 
 
def video_to_tensor(video_path):
    """
    Uniformly sample N_FRAMES frames (np.linspace), letterbox + ImageNet
    normalize. Returns tensor (1,16,3,224,224) float32, frame indices used,
    total_frames, fps, duration.
    """
    cap = cv2.VideoCapture(str(video_path))
    if not cap.isOpened():
        raise ValueError("Could not open the video file.")
 
    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    fps          = cap.get(cv2.CAP_PROP_FPS) or 25.0
    duration     = total_frames / fps if fps > 0 else 0.0
 
    if total_frames < 1:
        cap.release()
        raise ValueError("Video has no readable frames.")
 
    indices = np.linspace(0, total_frames - 1, N_FRAMES, dtype=int)
    indices = np.clip(indices, 0, total_frames - 1)
 
    frames = []
    for idx in indices:
        cap.set(cv2.CAP_PROP_POS_FRAMES, int(idx))
        ret, frame = cap.read()
        if not ret:
            fallback = frames[-1].copy() if frames else np.zeros(
                (IMG_SIZE, IMG_SIZE, 3), dtype=np.uint8)
            frames.append(fallback)
            continue
        frames.append(letterbox_frame(frame, IMG_SIZE))
    cap.release()
 
    arr = np.stack(frames, axis=0).astype(np.float32) / 255.0  # (16,224,224,3)
    arr = (arr - MEAN) / STD
    arr = arr.transpose(0, 3, 1, 2)                              # (16,3,224,224)
    tensor = torch.from_numpy(arr).unsqueeze(0)                  # (1,16,3,224,224)
 
    return tensor, indices, total_frames, fps, duration
 
 
 
# Inference (single GPU allocation)
 
@spaces.GPU(duration=60)
def run_inference_gpu(tensor):
    device = "cuda" if torch.cuda.is_available() else "cpu"
    MODEL.to(device)
    MODEL.eval()
 
    with torch.no_grad():
        t0 = time.time()
        x = tensor.to(device)
        logit, attn_weights = MODEL(x)
        prob = torch.sigmoid(logit).squeeze().item()
        elapsed_ms = (time.time() - t0) * 1000
 
    attn_np = attn_weights.squeeze(0).cpu().numpy()  # (16,)
    MODEL.to("cpu")
    return prob, attn_np, elapsed_ms
 
 
# Key frame extraction (peak attention frame)
 
def get_key_frame(video_path, indices, attn_weights, prob, fps):
    peak_pos   = int(np.argmax(attn_weights))
    global_idx = int(indices[peak_pos])
 
    cap = cv2.VideoCapture(str(video_path))
    cap.set(cv2.CAP_PROP_POS_FRAMES, global_idx)
    ret, frame = cap.read()
    cap.release()
    if not ret:
        return None
 
    t_sec = global_idx / fps if fps > 0 else 0.0
    label = "FALL" if prob >= FALL_THRESHOLD else "NORMAL"
    color = (0, 48, 255) if prob >= FALL_THRESHOLD else (0, 200, 100)  # BGR
 
    h, w = frame.shape[:2]
    cv2.rectangle(frame, (0, 0), (w, 42), (0, 0, 0), -1)
    cv2.putText(frame, f"PEAK ATTENTION  t={t_sec:.1f}s  p={prob:.3f}  {label}",
                (12, 28), cv2.FONT_HERSHEY_SIMPLEX, 0.65, color, 2, cv2.LINE_AA)
    cv2.rectangle(frame, (1, 1), (w - 2, h - 2), (0, 140, 255), 2)  # orange border (BGR)
 
    return cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
 
 
 
# Orange cyberpunk HUD report: attention bar chart
 
def build_attention_report(attn_weights, indices, fps, prob):
    plt.rcParams.update({
        "font.family": "monospace",
        "text.color": SILVER,
        "axes.edgecolor": ORANGE,
        "axes.labelcolor": SILVER,
        "xtick.color": SILVER,
        "ytick.color": SILVER,
    })
 
    timestamps = indices / fps if fps > 0 else indices
    peak_idx   = int(np.argmax(attn_weights))
 
    fig, ax = plt.subplots(figsize=(12, 5))
    fig.patch.set_facecolor(BG)
    ax.set_facecolor(PANEL)
    for s in ax.spines.values():
        s.set_color(ORANGE)
        s.set_linewidth(1.2)
    ax.grid(True, color=GRID, linewidth=0.6, axis="y")
 
    colors = [RED if i == peak_idx else ORANGE for i in range(len(attn_weights))]
    bars = ax.bar(range(len(attn_weights)), attn_weights, color=colors,
                  edgecolor=AMBER, linewidth=0.8)
 
    ax.set_xticks(range(len(attn_weights)))
    ax.set_xticklabels([f"{t:.1f}s" for t in timestamps], rotation=45, fontsize=8)
    ax.set_xlabel("Frame timestamp")
    ax.set_ylabel("Attention weight")
 
    verdict = "FALL" if prob >= FALL_THRESHOLD else "NORMAL"
    ax.set_title(
        f"AEGIS-SAFE-WORK // FALL DETECTOR — TEMPORAL ATTENTION\n"
        f"prob={prob:.4f}  threshold={FALL_THRESHOLD}  verdict={verdict}",
        color=AMBER, fontsize=11, loc="left"
    )
 
    ax.annotate("PEAK", xy=(peak_idx, attn_weights[peak_idx]),
                xytext=(peak_idx, attn_weights[peak_idx] + 0.03),
                color=RED, fontsize=9, fontweight="bold", ha="center")
 
    fig.tight_layout()
    buf = io.BytesIO()
    fig.savefig(buf, format="png", dpi=140, bbox_inches="tight", facecolor=BG)
    plt.close(fig)
    buf.seek(0)
    return Image.open(buf)
 
 
 
# Main handler
 
def analyze(video_path):
    empty_df = pd.DataFrame()
    if not video_path:
        return "Please upload a video.", None, None, empty_df
 
    cap = cv2.VideoCapture(str(video_path))
    fps_check    = cap.get(cv2.CAP_PROP_FPS) or 25.0
    total_check  = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    cap.release()
    duration_check = total_check / fps_check if fps_check > 0 else 0.0
    if duration_check > MAX_DURATION:
        return (f"Input Video cannot be longer than {int(MAX_DURATION)} sec "
                f"(got {duration_check:.1f}s).", None, None, empty_df)
 
    try:
        tensor, indices, total_frames, fps, duration = video_to_tensor(video_path)
    except ValueError as e:
        return str(e), None, None, empty_df
 
    prob, attn_weights, elapsed_ms = run_inference_gpu(tensor)
 
    report_img = build_attention_report(attn_weights, indices, fps, prob)
    key_frame  = get_key_frame(video_path, indices, attn_weights, prob, fps)
 
    verdict = "FALL" if prob >= FALL_THRESHOLD else "NORMAL"
    status  = f"Analysis complete  //  VERDICT: {verdict}  //  p={prob:.4f}"
 
    summary_df = pd.DataFrame({
        "Metric": [
            "Duration (s)", "FPS", "Frames sampled",
            "Fall probability", "Threshold", "Verdict",
            "Attention sum (should be ~1.0)", "Inference latency (ms)",
        ],
        "Value": [
            f"{duration:.1f}", f"{fps:.1f}", N_FRAMES,
            f"{prob:.4f}", f"{FALL_THRESHOLD}", verdict,
            f"{attn_weights.sum():.4f}", f"{elapsed_ms:.1f}",
        ],
    })
 
    return status, report_img, key_frame, summary_df
 
 
 
# Gradio UI - Orange Cyberpunk HUD
 
CSS = """
.gradio-container { background: #000000 !important; }
h1, h2, h3, p, span, label { color: #FFB347 !important; font-family: monospace !important; }
.block, .form { border: 1px solid #FF8C00 !important; border-radius: 6px !important;
                background: #0a0a0a !important; }
.gr-button { border: 1px solid #FF8C00 !important; color: #FF8C00 !important;
             background: #050505 !important; font-family: monospace !important; }
"""
 
with gr.Blocks(css=CSS, title="Aegis-Safe-Work Fall Detector") as demo:
    gr.Markdown("# AEGIS-SAFE-WORK // FALL DETECTOR")
    gr.Markdown(
        f"EfficientNet-Lite0 + Temporal Attention. Upload a short clip "
        f"(max {int(MAX_DURATION)} s) to run fall detection. The model samples "
        f"{N_FRAMES} frames uniformly across the clip and returns a single "
        f"verdict, together with the per-frame attention weights and the "
        f"peak-attention frame."
    )
 
    with gr.Row():
        with gr.Column(scale=1):
            video_in = gr.Video(label=f"Input video (<= {int(MAX_DURATION)} s)")
            run_btn  = gr.Button("RUN ANALYSIS", variant="primary")
            status   = gr.Markdown()
        with gr.Column(scale=1):
            key_out = gr.Image(label="Peak attention frame", type="numpy")
 
    report_out  = gr.Image(label="Temporal attention HUD", type="pil")
    summary_out = gr.Dataframe(label="Summary metrics", interactive=False)
 
    run_btn.click(
        fn=analyze,
        inputs=[video_in],
        outputs=[status, report_out, key_out, summary_out],
    )
 
 
if __name__ == "__main__":

    
    demo.queue().launch(show_api=False, ssr_mode=False)