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)