import spaces import cv2 import numpy as np from ai_edge_litert.interpreter import Interpreter as tflite_Interpreter from collections import deque import gradio as gr import os import sys import tempfile from ultralytics import YOLO # Fix harmless asyncio cleanup crash on Windows (WinError 10054) if sys.platform == "win32": import asyncio asyncio.set_event_loop_policy(asyncio.WindowsSelectorEventLoopPolicy()) # --- Configuration --- MODEL_PATH = "fall_detection_transformer.tflite" YOLO_POSE_MODEL = "yolo26n-pose.pt" INPUT_TIMESTEPS = 30 FALL_CONFIDENCE_THRESHOLD = 0.90 MIN_KEYPOINT_CONFIDENCE_FOR_NORMALIZATION = 0.3 FALL_EVENT_COOLDOWN = 10 # YOLO26-pose outputs 17 COCO keypoints in this order: YOLO_KEYPOINT_NAMES = [ "Nose", "Left Eye", "Right Eye", "Left Ear", "Right Ear", "Left Shoulder", "Right Shoulder", "Left Elbow", "Right Elbow", "Left Wrist", "Right Wrist", "Left Hip", "Right Hip", "Left Knee", "Right Knee", "Left Ankle", "Right Ankle", ] YOUR_KEYPOINT_NAMES_TRAINING = [ "Nose", "Left Eye", "Right Eye", "Left Ear", "Right Ear", "Left Shoulder", "Right Shoulder", "Left Elbow", "Right Elbow", "Left Wrist", "Right Wrist", "Left Hip", "Right Hip", "Left Knee", "Right Knee", "Left Ankle", "Right Ankle" ] SORTED_YOUR_KEYPOINT_NAMES = sorted(YOUR_KEYPOINT_NAMES_TRAINING) KEYPOINT_DICT_TRAINING = {name: i for i, name in enumerate(SORTED_YOUR_KEYPOINT_NAMES)} NUM_KEYPOINTS_TRAINING = len(KEYPOINT_DICT_TRAINING) NUM_FEATURES = NUM_KEYPOINTS_TRAINING * 3 YOLO_IDX_TO_NAME = {i: name for i, name in enumerate(YOLO_KEYPOINT_NAMES) if name in KEYPOINT_DICT_TRAINING} print("--- Initializing ---") print(f"NUM_FEATURES: {NUM_FEATURES}") # --- Load YOLO26 Pose Model --- print(f"Loading YOLO26 pose model: {YOLO_POSE_MODEL}") yolo_pose = YOLO(YOLO_POSE_MODEL) print("YOLO26 pose model loaded.") # --- Load TFLite Transformer Model --- interpreter = tflite_Interpreter(model_path=MODEL_PATH) interpreter.allocate_tensors() input_details = interpreter.get_input_details() output_details = interpreter.get_output_details() print(f"TFLite Model Loaded: {MODEL_PATH}") model_expected_shape = tuple(input_details[0]["shape"]) print(f"Model expected input shape: {model_expected_shape}") def get_kpt_indices_training_order(keypoint_name): kp_idx = KEYPOINT_DICT_TRAINING[keypoint_name] return kp_idx * 3, kp_idx * 3 + 1, kp_idx * 3 + 2 def normalize_skeleton_frame(frame_features_sorted, min_confidence=MIN_KEYPOINT_CONFIDENCE_FOR_NORMALIZATION): normalized_frame = np.copy(frame_features_sorted) try: ls_x_idx, ls_y_idx, ls_c_idx = get_kpt_indices_training_order("Left Shoulder") rs_x_idx, rs_y_idx, rs_c_idx = get_kpt_indices_training_order("Right Shoulder") lh_x_idx, lh_y_idx, lh_c_idx = get_kpt_indices_training_order("Left Hip") rh_x_idx, rh_y_idx, rh_c_idx = get_kpt_indices_training_order("Right Hip") except Exception as e: return frame_features_sorted ls_c = frame_features_sorted[ls_c_idx] rs_c = frame_features_sorted[rs_c_idx] lh_c = frame_features_sorted[lh_c_idx] rh_c = frame_features_sorted[rh_c_idx] ls_x, ls_y = frame_features_sorted[ls_x_idx], frame_features_sorted[ls_y_idx] rs_x, rs_y = frame_features_sorted[rs_x_idx], frame_features_sorted[rs_y_idx] lh_x, lh_y = frame_features_sorted[lh_x_idx], frame_features_sorted[lh_y_idx] rh_x, rh_y = frame_features_sorted[rh_x_idx], frame_features_sorted[rh_y_idx] mid_shoulder_x, mid_shoulder_y = np.nan, np.nan valid_ls = ls_c > min_confidence valid_rs = rs_c > min_confidence if valid_ls and valid_rs: mid_shoulder_x, mid_shoulder_y = (ls_x + rs_x) / 2, (ls_y + rs_y) / 2 elif valid_ls: mid_shoulder_x, mid_shoulder_y = ls_x, ls_y elif valid_rs: mid_shoulder_x, mid_shoulder_y = rs_x, rs_y mid_hip_x, mid_hip_y = np.nan, np.nan valid_lh = lh_c > min_confidence valid_rh = rh_c > min_confidence if valid_lh and valid_rh: mid_hip_x, mid_hip_y = (lh_x + rh_x) / 2, (lh_y + rh_y) / 2 elif valid_lh: mid_hip_x, mid_hip_y = lh_x, lh_y elif valid_rh: mid_hip_x, mid_hip_y = rh_x, rh_y ref_x, ref_y = np.nan, np.nan if not np.isnan(mid_shoulder_x) and not np.isnan(mid_hip_x): ref_x = (mid_shoulder_x + mid_hip_x) / 2 ref_y = (mid_shoulder_y + mid_hip_y) / 2 elif not np.isnan(mid_shoulder_x): ref_x, ref_y = mid_shoulder_x, mid_shoulder_y elif not np.isnan(mid_hip_x): ref_x, ref_y = mid_hip_x, mid_hip_y scale = 1.0 if not np.isnan(mid_shoulder_x) and not np.isnan(mid_hip_x): torso_height = np.sqrt((mid_shoulder_x - mid_hip_x)**2 + (mid_shoulder_y - mid_hip_y)**2) if torso_height > 1e-6: scale = torso_height if np.isnan(ref_x): return frame_features_sorted for kp_name in SORTED_YOUR_KEYPOINT_NAMES: x_idx, y_idx, c_idx = get_kpt_indices_training_order(kp_name) if frame_features_sorted[c_idx] > min_confidence: normalized_frame[x_idx] = (frame_features_sorted[x_idx] - ref_x) / scale normalized_frame[y_idx] = (frame_features_sorted[y_idx] - ref_y) / scale else: normalized_frame[x_idx] = 0.0 normalized_frame[y_idx] = 0.0 return normalized_frame def extract_features_from_yolo_result(result, frame_width, frame_height): frame_features = np.zeros(NUM_FEATURES, dtype=np.float32) if result.keypoints is None or len(result.keypoints) == 0: return frame_features if result.boxes is not None and len(result.boxes) > 0: areas = result.boxes.xywh[:, 2] * result.boxes.xywh[:, 3] best_idx = int(areas.argmax()) else: best_idx = 0 kps = result.keypoints[best_idx] kps_data = kps.data[0].cpu().numpy() # (17, 3) for yolo_idx, kp_name in YOLO_IDX_TO_NAME.items(): x_norm = kps_data[yolo_idx, 0] / frame_width y_norm = kps_data[yolo_idx, 1] / frame_height conf = float(kps_data[yolo_idx, 2]) x_idx, y_idx, c_idx = get_kpt_indices_training_order(kp_name) frame_features[x_idx] = x_norm frame_features[y_idx] = y_norm frame_features[c_idx] = conf frame_features = normalize_skeleton_frame(frame_features) return frame_features @spaces.GPU(duration=300) # allocate GPU for the full video processing window def run_fall_detection_on_video(video_path, progress=gr.Progress(track_tqdm=True)): if video_path is None: return None, "Please upload a video file." cap = cv2.VideoCapture(video_path) if not cap.isOpened(): return None, f"Error: Cannot open video: {video_path}" frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) fps = cap.get(cv2.CAP_PROP_FPS) or 30.0 total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) # Unique temp file per request — prevents concurrent users overwriting each other fd, output_path = tempfile.mkstemp(suffix=".mp4", prefix="fall_detection_", dir=tempfile.gettempdir()) os.close(fd) fourcc = cv2.VideoWriter_fourcc(*"avc1") # H.264 — playable in browsers out = cv2.VideoWriter(output_path, fourcc, fps, (frame_width, frame_height)) # Fallback: if avc1 is unavailable (OpenCV build without H.264), try mp4v if not out.isOpened(): fourcc = cv2.VideoWriter_fourcc(*"mp4v") out = cv2.VideoWriter(output_path, fourcc, fps, (frame_width, frame_height)) frame_buffer = deque(maxlen=INPUT_TIMESTEPS) fall_event_count = 0 last_fall_frame = -FALL_EVENT_COOLDOWN * fps frame_idx = 0 fall_detected_this_frame = False fall_confidence_val = 0.0 print(f"Processing video: {video_path} ({total_frames} frames, {fps:.1f} fps)") while True: ret, frame = cap.read() if not ret: break frame_idx += 1 fall_detected_this_frame = False results = yolo_pose(frame, verbose=False, device="cuda") result = results[0] features = extract_features_from_yolo_result(result, frame_width, frame_height) frame_buffer.append(features) if len(frame_buffer) == INPUT_TIMESTEPS: input_seq = np.array(frame_buffer, dtype=np.float32)[np.newaxis, ...] interpreter.set_tensor(input_details[0]["index"], input_seq) interpreter.invoke() output_data = interpreter.get_tensor(output_details[0]["index"]) fall_confidence_val = float(output_data[0][0]) if (fall_confidence_val >= FALL_CONFIDENCE_THRESHOLD and (frame_idx - last_fall_frame) >= FALL_EVENT_COOLDOWN * fps): fall_detected_this_frame = True fall_event_count += 1 last_fall_frame = frame_idx print(f" FALL DETECTED at frame {frame_idx} | confidence: {fall_confidence_val:.4f}") annotated_frame = result.plot(kpt_radius=4, line_width=2) status_text = "FALL DETECTED!" if fall_detected_this_frame else "Normal" conf_text = f"Conf: {fall_confidence_val:.2f}" if len(frame_buffer) == INPUT_TIMESTEPS else "Warming up..." overlay_color = (0, 0, 255) if fall_detected_this_frame else (0, 200, 0) cv2.rectangle(annotated_frame, (0, 0), (350, 70), (0, 0, 0), -1) cv2.putText(annotated_frame, status_text, (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.9, overlay_color, 2) cv2.putText(annotated_frame, conf_text, (10, 58), cv2.FONT_HERSHEY_SIMPLEX, 0.65, (255, 255, 255), 1) cv2.rectangle(annotated_frame, (frame_width - 200, 0), (frame_width, 40), (0, 0, 0), -1) cv2.putText(annotated_frame, f"Falls: {fall_event_count}", (frame_width - 190, 28), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 140, 255), 2) out.write(annotated_frame) cap.release() out.release() summary = ( f"Processing complete!\n" f"Total frames processed : {frame_idx}\n" f"Fall events detected : {fall_event_count}\n" f"Threshold used : {FALL_CONFIDENCE_THRESHOLD:.0%}\n" f"Model : YOLO26-pose + TFLite Transformer" ) print(summary) return output_path, summary EXAMPLE_VIDEOS = [f for f in [ "fall_example_1.mp4", "fall_example_2.mp4", "fall_example_3.mp4", "fall_example_4.mp4" ] if os.path.exists(f)] with gr.Blocks( title="Fall Detection - YOLO26 + Transformer", theme=gr.themes.Soft(primary_hue="red"), css="footer { display: none !important; }", ) as demo: gr.Markdown( """ # Fall Detection System ### YOLO26-pose keypoint extraction + TFLite Transformer classifier Upload a video. The system extracts 17-keypoint poses with YOLO26, feeds a 30-frame sliding window into the Transformer, and flags falls >= 90% confidence. """ ) with gr.Row(): with gr.Column(scale=1): video_input = gr.Video(label="Upload Video", sources=["upload"]) run_btn = gr.Button("Detect Falls", variant="primary", size="lg") gr.Markdown( f"**Pose Extractor:** YOLO26n-pose \n" f"**Classifier:** TFLite Transformer \n" f"**Sequence:** {INPUT_TIMESTEPS} frames \n" f"**Threshold:** {FALL_CONFIDENCE_THRESHOLD:.0%} \n" f"**Keypoints:** {NUM_KEYPOINTS_TRAINING} (COCO-17)" ) with gr.Column(scale=1): video_output = gr.Video(label="Annotated Output") summary_output = gr.Textbox(label="Detection Summary", lines=6, interactive=False) if EXAMPLE_VIDEOS: gr.Markdown("### Example Videos") gr.Examples(examples=[[v] for v in EXAMPLE_VIDEOS], inputs=video_input, label="Click to load") run_btn.click( fn=run_fall_detection_on_video, inputs=[video_input], outputs=[video_output, summary_output], show_progress=True ) if __name__ == "__main__": demo.launch()