Spaces:
Sleeping
Sleeping
| 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 | |
| # 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") 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( | |
| server_name="0.0.0.0", | |
| server_port=7860, | |
| share=False, | |
| show_error=True, | |
| ssr_mode=False, | |
| theme=gr.themes.Soft(primary_hue="red"), | |
| css="footer { display: none !important; }", | |
| ) | |