Spaces:
Runtime error
Runtime error
Update main.py
Browse files
main.py
CHANGED
@@ -1,30 +1,27 @@
|
|
1 |
import gradio as gr
|
2 |
from ultralytics import YOLO
|
3 |
-
import ai_gym
|
4 |
import cv2
|
5 |
import tempfile
|
6 |
-
import
|
7 |
|
8 |
-
|
9 |
-
model = YOLO("yolov8n-pose.pt")
|
10 |
-
|
11 |
-
# Initialize AIGym object
|
12 |
-
gym_object = ai_gym.AIGym()
|
13 |
-
|
14 |
-
def count_workouts(input_video):
|
15 |
-
# Temporary file to store output video
|
16 |
-
output_path = tempfile.NamedTemporaryFile(suffix='.avi').name
|
17 |
-
|
18 |
-
# Open input video
|
19 |
-
cap = cv2.VideoCapture(input_video.name)
|
20 |
assert cap.isOpened(), "Error reading video file"
|
21 |
w, h, fps = (int(cap.get(x)) for x in (cv2.CAP_PROP_FRAME_WIDTH, cv2.CAP_PROP_FRAME_HEIGHT, cv2.CAP_PROP_FPS))
|
22 |
|
23 |
-
#
|
24 |
-
video_writer = cv2.VideoWriter(
|
25 |
-
|
26 |
-
|
27 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
28 |
|
29 |
frame_count = 0
|
30 |
while cap.isOpened():
|
@@ -33,57 +30,25 @@ def count_workouts(input_video):
|
|
33 |
print("Video frame is empty or video processing has been successfully completed.")
|
34 |
break
|
35 |
frame_count += 1
|
36 |
-
|
37 |
-
|
38 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
39 |
|
40 |
cap.release()
|
41 |
-
video_writer.release()
|
42 |
cv2.destroyAllWindows()
|
43 |
|
44 |
-
return
|
45 |
-
|
46 |
-
# Gradio Interface
|
47 |
-
inputs = gr.inputs.Video(label="Upload a video")
|
48 |
-
outputs = gr.outputs.Video(label="Output Video")
|
49 |
-
|
50 |
-
gr.Interface(count_workouts, inputs, outputs, title="Workout Counter",
|
51 |
-
description="Upload a video and get a video with workout counting annotations.").launch()
|
52 |
-
|
53 |
-
|
54 |
-
|
55 |
-
# from ultralytics import YOLO
|
56 |
-
# from ultralytics.solutions import ai_gym
|
57 |
-
# import cv2
|
58 |
-
|
59 |
-
# model = YOLO("yolov8n-pose.pt")
|
60 |
-
# cap = cv2.VideoCapture("pullups.mp4")
|
61 |
-
# assert cap.isOpened(), "Error reading video file"
|
62 |
-
# w, h, fps = (int(cap.get(x)) for x in (cv2.CAP_PROP_FRAME_WIDTH, cv2.CAP_PROP_FRAME_HEIGHT, cv2.CAP_PROP_FPS))
|
63 |
-
|
64 |
-
# video_writer = cv2.VideoWriter("output_video.avi",
|
65 |
-
# cv2.VideoWriter_fourcc(*'mp4v'),
|
66 |
-
# fps,
|
67 |
-
# (w, h))
|
68 |
-
|
69 |
-
# gym_object = ai_gym.AIGym() # init AI GYM module
|
70 |
-
# gym_object.set_args(line_thickness=2,
|
71 |
-
# view_img=False, # Set view_img to False to prevent displaying the video in real-time
|
72 |
-
# pose_type="pushup",
|
73 |
-
# kpts_to_check=[6, 8, 10])
|
74 |
|
75 |
-
#
|
76 |
-
|
77 |
-
|
78 |
-
|
79 |
-
# print("Video frame is empty or video processing has been successfully completed.")
|
80 |
-
# break
|
81 |
-
# frame_count += 1
|
82 |
-
# results = model.track(im0, verbose=False) # Tracking recommended
|
83 |
-
# #results = model.predict(im0) # Prediction also supported
|
84 |
-
# im0 = gym_object.start_counting(im0, results, frame_count)
|
85 |
-
# video_writer.write(im0)
|
86 |
|
87 |
-
#
|
88 |
-
|
89 |
-
# cv2.destroyAllWindows()
|
|
|
1 |
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 |
|
8 |
+
def process(video_path):
|
9 |
+
model = YOLO("yolov8n-pose.pt")
|
10 |
+
cap = cv2.VideoCapture(video_path)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
11 |
assert cap.isOpened(), "Error reading video file"
|
12 |
w, h, fps = (int(cap.get(x)) for x in (cv2.CAP_PROP_FRAME_WIDTH, cv2.CAP_PROP_FRAME_HEIGHT, cv2.CAP_PROP_FPS))
|
13 |
|
14 |
+
temp_dir = tempfile.mkdtemp() # Create a temporary directory to store processed frames
|
15 |
+
video_writer = cv2.VideoWriter("output_video.mp4",
|
16 |
+
cv2.VideoWriter_fourcc(*'mp4v'),
|
17 |
+
fps,
|
18 |
+
(w, h))
|
19 |
+
|
20 |
+
gym_object = ai_gym.AIGym() # init AI GYM module
|
21 |
+
gym_object.set_args(line_thickness=2,
|
22 |
+
view_img=False, # Set view_img to False to prevent displaying the video in real-time
|
23 |
+
pose_type="pushup",
|
24 |
+
kpts_to_check=[6, 8, 10])
|
25 |
|
26 |
frame_count = 0
|
27 |
while cap.isOpened():
|
|
|
30 |
print("Video frame is empty or video processing has been successfully completed.")
|
31 |
break
|
32 |
frame_count += 1
|
33 |
+
if frame_count % 5 == 0: # Process every 5th frame
|
34 |
+
results = model.track(im0, verbose=False) # Tracking recommended
|
35 |
+
im0 = gym_object.start_counting(im0, results, frame_count)
|
36 |
+
# Save processed frame as an image in the temporary directory
|
37 |
+
cv2.imwrite(f"{temp_dir}/{frame_count}.jpg", im0)
|
38 |
+
|
39 |
+
# Use PIL to create the final video from the processed frames
|
40 |
+
images = [Image.open(f"{temp_dir}/{i}.jpg") for i in range(1, frame_count + 1)]
|
41 |
+
images[0].save("output_video.mp4", save_all=True, append_images=images[1:], duration=1000/fps, loop=0)
|
42 |
|
43 |
cap.release()
|
|
|
44 |
cv2.destroyAllWindows()
|
45 |
|
46 |
+
return "output_video.mp4"
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
47 |
|
48 |
+
# Create the Gradio demo
|
49 |
+
demo = gr.Interface(fn=process,
|
50 |
+
inputs=gr.Video(label='Input Video'),
|
51 |
+
outputs=gr.Video(label='Processed Video'))
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
52 |
|
53 |
+
# Launch the demo!
|
54 |
+
demo.launch(show_api=False)
|
|