Spaces:
Sleeping
Sleeping
| import streamlit as st | |
| import os | |
| import cv2 | |
| import numpy as np | |
| from ultralytics import YOLO | |
| # Load YOLO model | |
| model = YOLO('yolov8n.pt') # Ensure you have the correct model file | |
| def process_video(video_path): | |
| cap = cv2.VideoCapture(video_path) | |
| width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) | |
| height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) | |
| fps = int(cap.get(cv2.CAP_PROP_FPS)) | |
| # Create a video writer to save the output | |
| output_path = os.path.join(os.getcwd(), "output.mp4") | |
| fourcc = cv2.VideoWriter_fourcc(*"mp4v") | |
| out = cv2.VideoWriter(output_path, fourcc, fps, (width, height)) | |
| while cap.isOpened(): | |
| ret, frame = cap.read() | |
| if not ret: | |
| break | |
| results = model(frame) | |
| for result in results: | |
| for bbox in result.boxes: | |
| x1, y1, x2, y2 = map(int, bbox.xyxy[0]) | |
| confidence = float(bbox.conf) | |
| cls = int(bbox.cls) | |
| if cls == 0: # Assuming class 0 is 'person' | |
| w = x2 - x1 | |
| h = y2 - y1 | |
| if h < w: | |
| color = (0, 0, 255) # Red color for fall detected | |
| label = "Fall Detected" | |
| else: | |
| color = (0, 255, 0) # Green color for normal detection | |
| label = "Person" | |
| cv2.rectangle(frame, (x1, y1), (x2, y2), color, 2) | |
| cv2.putText(frame, label, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2) | |
| out.write(frame) | |
| cap.release() | |
| out.release() | |
| # Double-check that the file was saved | |
| if os.path.exists(output_path): | |
| st.success(f"Video successfully processed and saved to {output_path}") | |
| else: | |
| st.error("Failed to save the processed video.") | |
| return output_path | |
| # Streamlit interface | |
| st.title("Fall Detection App") | |
| st.write("The default video is automatically processed to detect falls.") | |
| # Create two columns | |
| left_column, right_column = st.columns(2) | |
| # Right column for video selection | |
| with right_column: | |
| default_videos = { | |
| "Video 1": os.path.join(os.getcwd(), "fall_test_01.mp4"), | |
| "Video 2": os.path.join(os.getcwd(), "fall_test_02.mp4"), | |
| "Video 3": "video3.mp4", | |
| } | |
| # Select the first video as the default input | |
| default_video_path = default_videos["Video 1"] | |
| # Display the selection to the user | |
| st.write("Default video selected: Video 1") | |
| # Process the default video automatically | |
| output_video = process_video(default_video_path) | |
| if output_video and os.path.exists(output_video): | |
| left_column.video(output_video) # Display video in the left column | |
| left_column.write("Download the processed video:") | |
| with open(output_video, "rb") as video_file: | |
| left_column.download_button("Download", video_file, "output.mp4") | |
| else: | |
| st.error("There was an issue processing the video. Please try again.") |