Spaces:
Runtime error
Runtime error
Update main.py
Browse files
main.py
CHANGED
@@ -2,23 +2,13 @@ import gradio as gr
|
|
2 |
from ultralytics import YOLO
|
3 |
from ultralytics.solutions import ai_gym
|
4 |
import cv2
|
5 |
-
import tempfile
|
6 |
-
from PIL import Image
|
7 |
-
import subprocess
|
8 |
|
9 |
-
# Function to upgrade pip
|
10 |
-
def upgrade_pip():
|
11 |
-
subprocess.run(['pip', 'install', '--upgrade', 'pip'])
|
12 |
-
|
13 |
-
# Function to process video
|
14 |
def process(video_path):
|
15 |
-
upgrade_pip() # Upgrade pip before executing the main function
|
16 |
model = YOLO("yolov8n-pose.pt")
|
17 |
cap = cv2.VideoCapture(video_path)
|
18 |
assert cap.isOpened(), "Error reading video file"
|
19 |
w, h, fps = (int(cap.get(x)) for x in (cv2.CAP_PROP_FRAME_WIDTH, cv2.CAP_PROP_FRAME_HEIGHT, cv2.CAP_PROP_FPS))
|
20 |
|
21 |
-
temp_dir = tempfile.mkdtemp() # Create a temporary directory to store processed frames
|
22 |
video_writer = cv2.VideoWriter("output_video.mp4",
|
23 |
cv2.VideoWriter_fourcc(*'mp4v'),
|
24 |
fps,
|
@@ -27,35 +17,109 @@ def process(video_path):
|
|
27 |
gym_object = ai_gym.AIGym() # init AI GYM module
|
28 |
gym_object.set_args(line_thickness=2,
|
29 |
view_img=False, # Set view_img to False to prevent displaying the video in real-time
|
30 |
-
pose_type="
|
31 |
kpts_to_check=[6, 8, 10])
|
32 |
|
33 |
frame_count = 0
|
34 |
while cap.isOpened():
|
35 |
success, im0 = cap.read()
|
36 |
if not success:
|
37 |
-
print("Video
|
38 |
break
|
39 |
frame_count += 1
|
40 |
-
|
41 |
-
|
42 |
-
|
43 |
-
# Save processed frame as an image in the temporary directory
|
44 |
-
cv2.imwrite(f"{temp_dir}/{frame_count}.jpg", im0)
|
45 |
-
|
46 |
-
# Use PIL to create the final video from the processed frames
|
47 |
-
images = [Image.open(f"{temp_dir}/{i}.jpg") for i in range(1, frame_count + 1)]
|
48 |
-
images[0].save("output_video.mp4", save_all=True, append_images=images[1:], duration=1000/fps, loop=0)
|
49 |
|
50 |
cap.release()
|
|
|
51 |
cv2.destroyAllWindows()
|
52 |
|
53 |
return "output_video.mp4"
|
54 |
|
|
|
|
|
|
|
|
|
|
|
|
|
55 |
# Create the Gradio demo
|
56 |
demo = gr.Interface(fn=process,
|
57 |
inputs=gr.Video(label='Input Video'),
|
58 |
-
outputs=gr.Video(label='
|
|
|
|
|
|
|
|
|
59 |
|
60 |
# Launch the demo!
|
61 |
-
demo.launch(show_api=
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
2 |
from ultralytics import YOLO
|
3 |
from ultralytics.solutions import ai_gym
|
4 |
import cv2
|
|
|
|
|
|
|
5 |
|
|
|
|
|
|
|
|
|
|
|
6 |
def process(video_path):
|
|
|
7 |
model = YOLO("yolov8n-pose.pt")
|
8 |
cap = cv2.VideoCapture(video_path)
|
9 |
assert cap.isOpened(), "Error reading video file"
|
10 |
w, h, fps = (int(cap.get(x)) for x in (cv2.CAP_PROP_FRAME_WIDTH, cv2.CAP_PROP_FRAME_HEIGHT, cv2.CAP_PROP_FPS))
|
11 |
|
|
|
12 |
video_writer = cv2.VideoWriter("output_video.mp4",
|
13 |
cv2.VideoWriter_fourcc(*'mp4v'),
|
14 |
fps,
|
|
|
17 |
gym_object = ai_gym.AIGym() # init AI GYM module
|
18 |
gym_object.set_args(line_thickness=2,
|
19 |
view_img=False, # Set view_img to False to prevent displaying the video in real-time
|
20 |
+
pose_type="pullup",
|
21 |
kpts_to_check=[6, 8, 10])
|
22 |
|
23 |
frame_count = 0
|
24 |
while cap.isOpened():
|
25 |
success, im0 = cap.read()
|
26 |
if not success:
|
27 |
+
print("Video processing has been successfully completed.")
|
28 |
break
|
29 |
frame_count += 1
|
30 |
+
results = model.track(im0, verbose=False) # Tracking recommended
|
31 |
+
im0 = gym_object.start_counting(im0, results, frame_count)
|
32 |
+
video_writer.write(im0)
|
|
|
|
|
|
|
|
|
|
|
|
|
33 |
|
34 |
cap.release()
|
35 |
+
video_writer.release()
|
36 |
cv2.destroyAllWindows()
|
37 |
|
38 |
return "output_video.mp4"
|
39 |
|
40 |
+
title = "Push-up Counter"
|
41 |
+
description = "This app counts the number of push-ups in a video."
|
42 |
+
# inputs = gr.inputs.Video(label='Input Video')
|
43 |
+
# outputs = gr.outputs.Video(label='Processed Video')
|
44 |
+
# example_list = ['pullups.mp4']
|
45 |
+
|
46 |
# Create the Gradio demo
|
47 |
demo = gr.Interface(fn=process,
|
48 |
inputs=gr.Video(label='Input Video'),
|
49 |
+
outputs=gr.Video(label='Output Video')
|
50 |
+
# title=title,
|
51 |
+
# description=description,
|
52 |
+
# examples=example_list
|
53 |
+
)
|
54 |
|
55 |
# Launch the demo!
|
56 |
+
demo.launch(show_api=True)
|
57 |
+
|
58 |
+
|
59 |
+
|
60 |
+
|
61 |
+
|
62 |
+
|
63 |
+
|
64 |
+
|
65 |
+
# import gradio as gr
|
66 |
+
# from ultralytics import YOLO
|
67 |
+
# from ultralytics.solutions import ai_gym
|
68 |
+
# import cv2
|
69 |
+
# import tempfile
|
70 |
+
# from PIL import Image
|
71 |
+
# import subprocess
|
72 |
+
|
73 |
+
# # Function to upgrade pip
|
74 |
+
# def upgrade_pip():
|
75 |
+
# subprocess.run(['pip', 'install', '--upgrade', 'pip'])
|
76 |
+
|
77 |
+
# # Function to process video
|
78 |
+
# def process(video_path):
|
79 |
+
# upgrade_pip() # Upgrade pip before executing the main function
|
80 |
+
# model = YOLO("yolov8n-pose.pt")
|
81 |
+
# cap = cv2.VideoCapture(video_path)
|
82 |
+
# assert cap.isOpened(), "Error reading video file"
|
83 |
+
# w, h, fps = (int(cap.get(x)) for x in (cv2.CAP_PROP_FRAME_WIDTH, cv2.CAP_PROP_FRAME_HEIGHT, cv2.CAP_PROP_FPS))
|
84 |
+
|
85 |
+
# temp_dir = tempfile.mkdtemp() # Create a temporary directory to store processed frames
|
86 |
+
# video_writer = cv2.VideoWriter("output_video.mp4",
|
87 |
+
# cv2.VideoWriter_fourcc(*'mp4v'),
|
88 |
+
# fps,
|
89 |
+
# (w, h))
|
90 |
+
|
91 |
+
# gym_object = ai_gym.AIGym() # init AI GYM module
|
92 |
+
# gym_object.set_args(line_thickness=2,
|
93 |
+
# view_img=False, # Set view_img to False to prevent displaying the video in real-time
|
94 |
+
# pose_type="pushup",
|
95 |
+
# kpts_to_check=[6, 8, 10])
|
96 |
+
|
97 |
+
# frame_count = 0
|
98 |
+
# while cap.isOpened():
|
99 |
+
# success, im0 = cap.read()
|
100 |
+
# if not success:
|
101 |
+
# print("Video frame is empty or video processing has been successfully completed.")
|
102 |
+
# break
|
103 |
+
# frame_count += 1
|
104 |
+
# if frame_count % 5 == 0: # Process every 5th frame
|
105 |
+
# results = model.track(im0, verbose=False) # Tracking recommended
|
106 |
+
# im0 = gym_object.start_counting(im0, results, frame_count)
|
107 |
+
# # Save processed frame as an image in the temporary directory
|
108 |
+
# cv2.imwrite(f"{temp_dir}/{frame_count}.jpg", im0)
|
109 |
+
|
110 |
+
# # Use PIL to create the final video from the processed frames
|
111 |
+
# images = [Image.open(f"{temp_dir}/{i}.jpg") for i in range(1, frame_count + 1)]
|
112 |
+
# images[0].save("output_video.mp4", save_all=True, append_images=images[1:], duration=1000/fps, loop=0)
|
113 |
+
|
114 |
+
# cap.release()
|
115 |
+
# cv2.destroyAllWindows()
|
116 |
+
|
117 |
+
# return "output_video.mp4"
|
118 |
+
|
119 |
+
# # Create the Gradio demo
|
120 |
+
# demo = gr.Interface(fn=process,
|
121 |
+
# inputs=gr.Video(label='Input Video'),
|
122 |
+
# outputs=gr.Video(label='Processed Video'))
|
123 |
+
|
124 |
+
# # Launch the demo!
|
125 |
+
# demo.launch(show_api=False)
|