Spaces:
Sleeping
Sleeping
File size: 2,213 Bytes
0db98b9 77ece12 0db98b9 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 |
import os
import numpy as np
import gradio as gr
import supervision as sv
from ultralytics import YOLO
# Define paths
HOME = os.getcwd()
MODEL_PATH = "./best.pt"
# Load the YOLO model
model = YOLO(MODEL_PATH)
# Initialize annotators
box_annotator = sv.BoxAnnotator()
label_annotator = sv.LabelAnnotator()
# Define the confidence threshold
CONFIDENCE_THRESHOLD = 0.6
# Define the callback function for processing each video frame
def callback(frame: np.ndarray, _: int) -> np.ndarray:
# Perform detection on the frame
results = model(frame)[0]
detections = sv.Detections.from_ultralytics(results)
# Filter detections based on confidence threshold
detections_filtered = detections[detections.confidence >
CONFIDENCE_THRESHOLD]
# Create labels for filtered detections
labels = [
f"{model.model.names[class_id]} {confidence:.2f}"
for class_id, confidence in zip(
detections_filtered.class_id, detections_filtered.confidence
)
]
# Annotate the frame with bounding boxes and labels
annotated_frame = box_annotator.annotate(
scene=frame.copy(),
detections=detections_filtered,
)
annotated_frame = label_annotator.annotate(
scene=annotated_frame,
detections=detections_filtered,
labels=labels
)
return annotated_frame
# Function to process the video and generate the output
def process_video_gradio(input_video):
SOURCE_VIDEO_PATH = input_video
TARGET_VIDEO_PATH = f"{HOME}/output_fall_detection.mp4"
sv.process_video(
source_path=SOURCE_VIDEO_PATH,
target_path=TARGET_VIDEO_PATH,
callback=callback
)
return TARGET_VIDEO_PATH
# Define the Gradio interface
interface = gr.Interface(
fn=process_video_gradio, # Function to process video
inputs=gr.Video(), # Upload video input
outputs=gr.Video(), # Return the annotated video output
title="Fall Detection Video Annotator",
description="Upload a video, and the model will annotate it with fall detection using Fine-Tuned YOLO model."
)
# Launch the interface
if __name__ == "__main__":
interface.launch()
|