Rename test.py to main.py

test.py → main.py

@@ -1,5 +1,59 @@
+import gradio as gr
+from ultralytics import YOLO
+from ultralytics.solutions import ai_gym
+import cv2
+import tempfile
+import os
+
+# Initialize YOLO model
+model = YOLO("yolov8n-pose.pt")
+
+# Initialize AIGym object
+gym_object = ai_gym.AIGym()
+
+def count_workouts(input_video):
+    # Temporary file to store output video
+    output_path = tempfile.NamedTemporaryFile(suffix='.avi').name
+
+    # Open input video
+    cap = cv2.VideoCapture(input_video.name)
+    assert cap.isOpened(), "Error reading video file"
+    w, h, fps = (int(cap.get(x)) for x in (cv2.CAP_PROP_FRAME_WIDTH, cv2.CAP_PROP_FRAME_HEIGHT, cv2.CAP_PROP_FPS))
+
+    # Initialize video writer for output video
+    video_writer = cv2.VideoWriter(output_path,
+                                    cv2.VideoWriter_fourcc(*'mp4v'),
+                                    fps,
+                                    (w, h))
+
+    frame_count = 0
+    while cap.isOpened():
+        success, im0 = cap.read()
+        if not success:
+            print("Video frame is empty or video processing has been successfully completed.")
+            break
+        frame_count += 1
+        results = model.track(im0, verbose=False)  # Tracking recommended
+        im0 = gym_object.start_counting(im0, results, frame_count)
+        video_writer.write(im0)
+
+    cap.release()
+    video_writer.release()
+    cv2.destroyAllWindows()
+
+    return output_path
+
+# Gradio Interface
+inputs = gr.inputs.Video(label="Upload a video")
+outputs = gr.outputs.Video(label="Output Video")
+
+gr.Interface(count_workouts, inputs, outputs, title="Workout Counter",
+             description="Upload a video and get a video with workout counting annotations.").launch()
+
+
+
 # from ultralytics import YOLO
-# import ai_gym
+# from ultralytics.solutions import ai_gym
 # import cv2
 
 # model = YOLO("yolov8n-pose.pt")
@@ -7,65 +61,29 @@
 # assert cap.isOpened(), "Error reading video file"
 # w, h, fps = (int(cap.get(x)) for x in (cv2.CAP_PROP_FRAME_WIDTH, cv2.CAP_PROP_FRAME_HEIGHT, cv2.CAP_PROP_FPS))
 
-# video_writer = cv2.VideoWriter("workouts.avi",
-#                                 cv2.VideoWriter_fourcc(*'mp4v'),
-#                                 fps,
-#                                 (w, h))
+# video_writer = cv2.VideoWriter("output_video.avi",
+#                                cv2.VideoWriter_fourcc(*'mp4v'),
+#                                fps,
+#                                (w, h))
 
 # gym_object = ai_gym.AIGym()  # init AI GYM module
 # gym_object.set_args(line_thickness=2,
-#                     view_img=True,
-#                     pose_type="pullup",
+#                     view_img=False,  # Set view_img to False to prevent displaying the video in real-time
+#                     pose_type="pushup",
 #                     kpts_to_check=[6, 8, 10])
 
 # frame_count = 0
 # while cap.isOpened():
 #     success, im0 = cap.read()
 #     if not success:
-#       print("Video frame is empty or video processing has been successfully completed.")
-#       break
+#         print("Video frame is empty or video processing has been successfully completed.")
+#         break
 #     frame_count += 1
 #     results = model.track(im0, verbose=False)  # Tracking recommended
 #     #results = model.predict(im0)  # Prediction also supported
 #     im0 = gym_object.start_counting(im0, results, frame_count)
 #     video_writer.write(im0)
 
-# cv2.destroyAllWindows()
+# cap.release()
 # video_writer.release()
-
-
-from ultralytics import YOLO
-from ultralytics.solutions import ai_gym
-import cv2
-
-model = YOLO("yolov8n-pose.pt")
-cap = cv2.VideoCapture("pullups.mp4")
-assert cap.isOpened(), "Error reading video file"
-w, h, fps = (int(cap.get(x)) for x in (cv2.CAP_PROP_FRAME_WIDTH, cv2.CAP_PROP_FRAME_HEIGHT, cv2.CAP_PROP_FPS))
-
-video_writer = cv2.VideoWriter("output_video.avi",
-                               cv2.VideoWriter_fourcc(*'mp4v'),
-                               fps,
-                               (w, h))
-
-gym_object = ai_gym.AIGym()  # init AI GYM module
-gym_object.set_args(line_thickness=2,
-                    view_img=False,  # Set view_img to False to prevent displaying the video in real-time
-                    pose_type="pushup",
-                    kpts_to_check=[6, 8, 10])
-
-frame_count = 0
-while cap.isOpened():
-    success, im0 = cap.read()
-    if not success:
-        print("Video frame is empty or video processing has been successfully completed.")
-        break
-    frame_count += 1
-    results = model.track(im0, verbose=False)  # Tracking recommended
-    #results = model.predict(im0)  # Prediction also supported
-    im0 = gym_object.start_counting(im0, results, frame_count)
-    video_writer.write(im0)
-
-cap.release()
-video_writer.release()
-cv2.destroyAllWindows()
+# cv2.destroyAllWindows()