Update main.py

main.py

@@ -2,23 +2,13 @@ import gradio as gr
 from ultralytics import YOLO
 import ai_gym
 import cv2
-import tempfile
-from PIL import Image
-import subprocess
 
-# Function to upgrade pip
-def upgrade_pip():
-    subprocess.run(['pip', 'install', '--upgrade', 'pip'])
-
-# Process video function
 def process(video_path, pose_type):
-    upgrade_pip()  # Upgrade pip before executing the main function
     model = YOLO("yolov8n-pose.pt")
     cap = cv2.VideoCapture(video_path)
     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))
 
-    temp_dir = tempfile.mkdtemp()  # Create a temporary directory to store processed frames
     video_writer = cv2.VideoWriter("output_video.mp4",
                                    cv2.VideoWriter_fourcc(*'mp4v'),
                                    fps,
@@ -37,17 +27,12 @@ def process(video_path, pose_type):
             print("Video processing has been successfully completed.")
             break
         frame_count += 1
-        if frame_count % 5 == 0:  # Process every 5th frame
-            results = model.track(im0, verbose=False)  # Tracking recommended
-            im0 = gym_object.start_counting(im0, results, frame_count)
-            # Save processed frame as an image in the temporary directory
-            cv2.imwrite(f"{temp_dir}/{frame_count}.jpg", im0)
-
-    # Use PIL to create the final video from the processed frames
-    images = [Image.open(f"{temp_dir}/{i}.jpg") for i in range(1, frame_count + 1)]
-    images[0].save("output_video.mp4", save_all=True, append_images=images[1:], duration=1000/fps, loop=0)
+        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_video.mp4"
@@ -57,17 +42,17 @@ description = "This app counts the number of push-ups in a video."
 inputs = [gr.Video(label='Input Video'),  
           gr.Radio(["pullups", "pushups", "absworkout"], label="Pose Type")]
 outputs = gr.Video(label='Output Video')
-example_list = [['Examples/PULL-UPS.mp4'],['Examples/PUSH-UPS.mp4']]
 
+example_list = [['Examples/PULL-UPS.mp4'],['Examples/PUSH-UPS.mp4']]
 # Create the Gradio demo
 demo = gr.Interface(fn=process,
                     inputs=inputs,
                     outputs=outputs,
                     title=title,
-                    description=description,
+                    description=description, 
                     examples=example_list,
-                    cache_examples=True,
-                    )
+                    cache_examples=True
+                   )
 
 # Launch the demo!
 demo.launch(show_api=True)