randomshit11 commited on
Commit
924c3e9
·
verified ·
1 Parent(s): 78c5956

Update main.py

Browse files
Files changed (1) hide show
  1. main.py +32 -67
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 os
7
 
8
- # Initialize YOLO model
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
- # Initialize video writer for output video
24
- video_writer = cv2.VideoWriter(output_path,
25
- cv2.VideoWriter_fourcc(*'mp4v'),
26
- fps,
27
- (w, h))
 
 
 
 
 
 
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
- results = model.track(im0, verbose=False) # Tracking recommended
37
- im0 = gym_object.start_counting(im0, results, frame_count)
38
- video_writer.write(im0)
 
 
 
 
 
 
39
 
40
  cap.release()
41
- video_writer.release()
42
  cv2.destroyAllWindows()
43
 
44
- return output_path
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
- # frame_count = 0
76
- # while cap.isOpened():
77
- # success, im0 = cap.read()
78
- # if not success:
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
- # cap.release()
88
- # video_writer.release()
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)