tstone87 commited on
Commit
7c90da7
·
verified ·
1 Parent(s): 53f61b7

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +58 -0
app.py ADDED
@@ -0,0 +1,58 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ import os
3
+ from ultralytics import YOLO
4
+
5
+ # Load the YOLO11-pose model (auto-downloads if not present)
6
+ model = YOLO("yolo11n-pose.pt")
7
+
8
+ def process_input(uploaded_file, youtube_link):
9
+ """
10
+ Process an uploaded file or a YouTube link to perform pose detection.
11
+ Returns the path to the annotated output.
12
+ """
13
+ if youtube_link and youtube_link.strip():
14
+ try:
15
+ from pytube import YouTube
16
+ yt = YouTube(youtube_link)
17
+ stream = yt.streams.filter(file_extension='mp4', progressive=True)\
18
+ .order_by("resolution").desc().first()
19
+ if stream is None:
20
+ return "No suitable mp4 stream found."
21
+ input_path = stream.download()
22
+ except Exception as e:
23
+ return f"Error downloading video: {e}"
24
+ elif uploaded_file is not None:
25
+ input_path = uploaded_file.name
26
+ else:
27
+ return "Please provide an uploaded file or a YouTube link."
28
+
29
+ # Run pose detection and save the annotated output.
30
+ results = model.predict(source=input_path, save=True)
31
+
32
+ try:
33
+ output_path = results[0].save_path
34
+ except Exception as e:
35
+ return f"Error processing the file: {e}"
36
+
37
+ # Optionally remove the downloaded video if applicable.
38
+ if youtube_link and os.path.exists(input_path):
39
+ os.remove(input_path)
40
+
41
+ return output_path
42
+
43
+ # Define the Gradio Blocks interface as a global variable.
44
+ demo = gr.Blocks()
45
+
46
+ with demo:
47
+ gr.Markdown("# Pose Detection with YOLO11-pose")
48
+ gr.Markdown("Upload an image/video or provide a YouTube link to detect human poses.")
49
+ with gr.Row():
50
+ file_input = gr.File(label="Upload Image/Video")
51
+ youtube_input = gr.Textbox(label="Or enter a YouTube link", placeholder="https://...")
52
+ output_file = gr.File(label="Download Annotated Output")
53
+ run_button = gr.Button("Run Pose Detection")
54
+ run_button.click(process_input, inputs=[file_input, youtube_input], outputs=output_file)
55
+
56
+ # Only launch the interface if this file is executed directly.
57
+ if __name__ == "__main__":
58
+ demo.launch()