import gradio as gr
from ultralytics import YOLO
import tempfile
import os
import cv2
import numpy as np
import torch
import atexit
import uuid
# Load the YOLOv8 pose estimation model once at the start
model = YOLO("yolov8n-pose.pt")
# Define the skeleton connections based on COCO keypoints
COCO_KEYPOINTS = [
"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"
]
# Define the skeleton as pairs of keypoints indices
SKELETON_CONNECTIONS = [
(0, 1), (0, 2), # Nose to eyes
(1, 3), (2, 4), # Eyes to ears
(0, 5), (0, 6), # Nose to shoulders
(5, 6), # Shoulders to each other
(5, 7), (6, 8), # Shoulders to elbows
(7, 9), (8, 10), # Elbows to wrists
(5, 11), (6, 12), # Shoulders to hips
(11, 12), # Hips to each other
(11, 13), (12, 14), # Hips to knees
(13, 15), (14, 16) # Knees to ankles
]
def calculate_torso_angle(keypoints, frame_height):
"""
Calculate the angle of the torso with respect to the vertical axis.
Args:
keypoints (numpy.ndarray): Array of shape (17, 3) representing COCO keypoints.
frame_height (int): Height of the video frame in pixels.
Returns:
float: Angle in degrees. Returns None if keypoints are not detected properly.
"""
try:
# COCO keypoint indices
LEFT_SHOULDER = 5
RIGHT_SHOULDER = 6
LEFT_HIP = 11
RIGHT_HIP = 12
# Extract shoulder and hip coordinates
left_shoulder = keypoints[LEFT_SHOULDER][:2]
right_shoulder = keypoints[RIGHT_SHOULDER][:2]
left_hip = keypoints[LEFT_HIP][:2]
right_hip = keypoints[RIGHT_HIP][:2]
# Check visibility (visibility > 0.3)
if (keypoints[LEFT_SHOULDER][2] < 0.3 or keypoints[RIGHT_SHOULDER][2] < 0.3 or
keypoints[LEFT_HIP][2] < 0.3 or keypoints[RIGHT_HIP][2] < 0.3):
return None
# Calculate mid points
mid_shoulder = (left_shoulder + right_shoulder) / 2
mid_hip = (left_hip + right_hip) / 2
# Calculate the vector of the torso
vector = mid_hip - mid_shoulder
# Calculate angle with respect to the vertical axis
angle_rad = np.arctan2(vector[0], vector[1])
angle_deg = np.degrees(angle_rad)
return angle_deg
except Exception as e:
print(f"Error calculating torso angle: {e}")
return None
def draw_skeleton(frame, keypoints, show_labels=True):
"""
Draws the skeleton on the frame based on keypoints.
Args:
frame (numpy.ndarray): The current video frame.
keypoints (numpy.ndarray): Array of shape (17, 3) representing COCO keypoints.
show_labels (bool): Whether to display keypoint indices.
Returns:
numpy.ndarray: Annotated frame with skeleton.
"""
for connection in SKELETON_CONNECTIONS:
start_idx, end_idx = connection
x_start, y_start, conf_start = keypoints[start_idx]
x_end, y_end, conf_end = keypoints[end_idx]
# Only draw if both keypoints have sufficient confidence
if conf_start > 0.5 and conf_end > 0.5:
start_point = (int(x_start), int(y_start))
end_point = (int(x_end), int(y_end))
cv2.line(frame, start_point, end_point, (255, 0, 0), 2) # Blue lines
if show_labels:
# Draw keypoints indices
for idx, (x, y, conf) in enumerate(keypoints):
if conf > 0.5:
cv2.putText(frame, f"{idx}", (int(x), int(y)), cv2.FONT_HERSHEY_SIMPLEX, 0.3, (255, 0, 0), 1) # Blue labels
return frame
def detect_fall(video_path, angle_threshold=30, consecutive_frames=3, frame_sampling_rate=1, confidence_threshold=0.3, show_labels=True):
"""
Detects falls in the uploaded video using pose estimation.
Args:
video_path (str): The path to the input video file uploaded by the user.
angle_threshold (float): Angle threshold to classify a fall (in degrees).
consecutive_frames (int): Number of consecutive frames to confirm a fall.
frame_sampling_rate (int): Process every nth frame.
confidence_threshold (float): Minimum confidence required for keypoint detection.
show_labels (bool): Whether to display keypoint indices.
Returns:
tuple: (annotated_video_path, notification_message)
"""
try:
cap = cv2.VideoCapture(video_path)
if not cap.isOpened():
raise ValueError("Unable to open the video file.")
# Video properties
fps = cap.get(cv2.CAP_PROP_FPS)
width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
fourcc = cv2.VideoWriter_fourcc(*'mp4v')
# Create a unique temporary file for the annotated video
unique_id = uuid.uuid4().hex
with tempfile.NamedTemporaryFile(delete=False, suffix=".mp4", prefix=f"annotated_{unique_id}_") as tmp:
annotated_video_path = tmp.name
out = cv2.VideoWriter(annotated_video_path, fourcc, fps, (width, height))
frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
current_frame = 0
consecutive_fall_frames = 0
total_falls = 0
fall_frames = [] # To store frames where falls were detected
while True:
ret, frame = cap.read()
if not ret:
break # End of video
current_frame += 1
# Implement frame sampling
if current_frame % frame_sampling_rate != 0:
out.write(frame)
continue
print(f"Processing frame {current_frame}/{frame_count}")
# Run pose estimation
results = model.predict(source=frame, conf=confidence_threshold, save=False, stream=False)
# Iterate through detected persons
for result in results:
if not hasattr(result, 'keypoints') or result.keypoints is None:
continue
for keypoints in result.keypoints.data:
# keypoints should be a tensor of shape (17,3)
if keypoints is None or not hasattr(keypoints, 'cpu'):
continue
# Convert to NumPy array
if isinstance(keypoints, torch.Tensor):
kpts = keypoints.cpu().numpy()
elif isinstance(keypoints, np.ndarray):
kpts = keypoints
else:
print(f"Unexpected keypoints data type: {type(keypoints)}")
continue
if kpts.size == 0 or kpts.shape[0] < 17:
print(f"Insufficient keypoints for processing in frame {current_frame}")
continue
angle = calculate_torso_angle(kpts, height)
if angle is None:
continue
# Determine if it's a fall
if abs(angle) > angle_threshold:
consecutive_fall_frames += 1
label = "Fall Detected!"
color = (0, 0, 255) # Red
else:
if consecutive_fall_frames >= consecutive_frames:
total_falls += 1
fall_frames.append(current_frame)
consecutive_fall_frames = 0
label = "Normal"
color = (0, 255, 0) # Green
# If fall persists over consecutive frames, mark as fall
if consecutive_fall_frames >= consecutive_frames:
cv2.putText(frame, label, (50, 50), cv2.FONT_HERSHEY_SIMPLEX, 1, color, 2)
# Draw keypoints and skeleton
frame = draw_skeleton(frame, kpts, show_labels=show_labels)
# Write annotated frame
out.write(frame)
# Release resources
cap.release()
out.release()
# Final check for falls that persisted until the end of the video
if consecutive_fall_frames >= consecutive_frames:
total_falls += 1
fall_frames.append(current_frame)
# Generate notification message
if total_falls > 0:
if total_falls == 1:
notification = f"A fall was detected at frame {fall_frames[0]}."
else:
frames = ', '.join(map(str, fall_frames))
notification = f"{total_falls} falls were detected at frames: {frames}."
else:
notification = "No falls were detected in the video."
# Check if annotated video was created
if not os.path.exists(annotated_video_path):
raise FileNotFoundError("Annotated video was not found. Please check the model and processing steps.")
return annotated_video_path, notification
except Exception as e:
# Clean up in case of an error
print(f"Error during fall detection: {e}")
return None, f"An error occurred during fall detection: {e}"
def create_gradio_interface():
# Define the Gradio interface with adjustable parameters
iface = gr.Interface(
fn=detect_fall,
inputs=[
gr.Video(label="Upload Video"),
gr.Slider(
label="Angle Threshold (degrees)",
minimum=0,
maximum=90,
step=1,
value=30,
interactive=True,
info="Adjust the torso angle threshold to classify a fall. Lower values increase sensitivity."
),
gr.Slider(
label="Consecutive Frames to Confirm Fall",
minimum=1,
maximum=10,
step=1,
value=3,
interactive=True,
info="Number of consecutive frames exceeding the angle threshold required to confirm a fall."
),
gr.Slider(
label="Frame Sampling Rate",
minimum=1,
maximum=10,
step=1,
value=1,
interactive=True,
info="Process every nth frame to speed up detection. Higher values reduce processing time."
),
gr.Slider(
label="Confidence Threshold",
minimum=0.0,
maximum=1.0,
step=0.05,
value=0.3, # Changed default value to 0.3
interactive=True,
info="Minimum confidence required for keypoint detection. Higher values reduce false positives."
),
gr.Checkbox(
label="Show Keypoint Labels",
value=True,
interactive=True,
info="Toggle the display of keypoint indices on the video."
)
],
outputs=[
gr.Video(label="Annotated Video"),
gr.Textbox(label="Fall Detection Notification")
],
title="Fall Detection App 🚨",
description=(
"Upload a video of a person falling, and the app will detect and annotate the fall "
"using pose estimation. Adjust the angle threshold, consecutive frames, frame sampling rate, "
"and confidence threshold to fine-tune detection sensitivity and performance. "
"The annotated video will display keypoints, skeleton lines, and indicate when a fall is detected."
),
examples=[
["demo/person falling.mp4", 30, 3, 1, 0.3, True]
], # Added example video with corresponding parameter values
flagging_mode="never", # Updated parameter name
)
return iface
# Create the Gradio interface
iface = create_gradio_interface()
# Ensure temporary directories are cleaned up on exit
def cleanup_temp_dirs():
temp_dir = tempfile.gettempdir()
# Implement additional cleanup logic if necessary
atexit.register(cleanup_temp_dirs)
# Launch the app
if __name__ == "__main__":
iface.launch()