randomshit11 commited on
Commit
40e0d6d
·
verified ·
1 Parent(s): d789906

Update main.py

Browse files
Files changed (1) hide show
  1. main.py +87 -23
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="pushup",
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 frame is empty or video processing has been successfully completed.")
38
  break
39
  frame_count += 1
40
- if frame_count % 5 == 0: # Process every 5th frame
41
- results = model.track(im0, verbose=False) # Tracking recommended
42
- im0 = gym_object.start_counting(im0, results, frame_count)
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='Processed Video'))
 
 
 
 
59
 
60
  # Launch the demo!
61
- demo.launch(show_api=False)
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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)