Update app.py

app.py

@@ -1,16 +1,76 @@
-# Import the InferencePipeline object
-from inference import InferencePipeline
-
-# Import the built in render_boxes sink for visualizing results
-from inference.core.interfaces.stream.sinks import render_boxes
-
-# initialize a pipeline object
-pipeline = InferencePipeline.init(
-    model_id="fall-detection-v2-ihywg/5",  # Roboflow model to use
-    # Path to video, device id (int, usually 0 for built in webcams), or RTSP stream url
-    video_reference="./examples/fall_test.mp4",
-    on_prediction=render_boxes,  # Function to run after each prediction
-    api_key="rOZmS39JRmLDRuDaF8KE"
+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."
 )
-pipeline.start()
-pipeline.join()
+
+# Launch the interface
+if __name__ == "__main__":
+    interface.launch()