Upload app.py

app.py

@@ -3,6 +3,7 @@ import cv2
 import numpy as np
 import tempfile
 import os
+import time
 from pose_detector import PoseDetector
 from dance_generator import DanceGenerator
 from dance_visualizer import DanceVisualizer
@@ -32,13 +33,21 @@ class AIDancePartner:
         # Add playback controls
         play_speed = st.sidebar.slider("Playback Speed", 0.1, 2.0, 1.0, 0.1)
         
+        # Add play/pause button
+        is_playing = st.sidebar.button("Play/Pause")
+        
         if video_file:
-            self.process_video(video_file, mode, play_speed)
+            self.process_video(video_file, mode, play_speed, is_playing)
 
-    def process_video(self, video_file, mode, play_speed):
+    def process_video(self, video_file, mode, play_speed, is_playing):
+        # Create a placeholder for the video
+        video_placeholder = st.empty()
+        
+        # Create temporary file
         tfile = tempfile.NamedTemporaryFile(delete=False)
         tfile.write(video_file.read())
         
+        # Open video file
         cap = cv2.VideoCapture(tfile.name)
         
         # Get video properties
@@ -49,59 +58,72 @@ class AIDancePartner:
         
         # Initialize progress bar
         progress_bar = st.progress(0)
-        frame_placeholder = st.empty()
-        
-        # Pre-process video to extract all poses
-        all_poses = []
-        while cap.isOpened():
-            ret, frame = cap.read()
-            if not ret:
-                break
-            pose_landmarks = self.pose_detector.detect_pose(frame)
-            all_poses.append(pose_landmarks)
         
-        # Reset video capture
-        cap.set(cv2.CAP_PROP_POS_FRAMES, 0)
+        # Store video frames and AI sequences
+        frames = []
+        ai_sequences = []
         
-        # Generate AI dance sequence
-        ai_sequence = self.dance_generator.generate_dance_sequence(
-            all_poses, 
-            mode, 
-            total_frames,
-            (frame_height, frame_width)
-        )
+        # Pre-process video
+        with st.spinner('Processing video...'):
+            while cap.isOpened():
+                ret, frame = cap.read()
+                if not ret:
+                    break
+                    
+                pose_landmarks = self.pose_detector.detect_pose(frame)
+                ai_frame = self.dance_generator.generate_dance_sequence(
+                    [pose_landmarks], 
+                    mode, 
+                    1,
+                    (frame_height, frame_width)
+                )[0]
+                
+                frames.append(frame)
+                ai_sequences.append(ai_frame)
         
         # Playback loop
         frame_count = 0
-        while cap.isOpened():
-            ret, frame = cap.read()
-            if not ret:
+        play = True if is_playing else False
+        
+        while True:
+            if play:
+                if frame_count >= len(frames):
+                    frame_count = 0
+                    
+                # Get current frames
+                frame = frames[frame_count]
+                ai_frame = ai_sequences[frame_count]
+                
+                # Combine frames side by side
+                combined_frame = np.hstack([
+                    frame,
+                    cv2.resize(ai_frame, (frame_width, frame_height))
+                ])
+                
+                # Update progress
+                progress = frame_count / total_frames
+                progress_bar.progress(progress)
+                
+                # Display frame
+                video_placeholder.image(
+                    combined_frame,
+                    channels="BGR",
+                    use_container_width=True
+                )
+                
+                # Control playback speed
+                time.sleep(1 / (fps * play_speed))
+                frame_count += 1
+                
+            # Check for play/pause button
+            if st.sidebar.button("Stop"):
                 break
                 
-            # Update progress
-            progress = frame_count / total_frames
-            progress_bar.progress(progress)
-            
-            # Get corresponding AI frame
-            ai_frame = ai_sequence[frame_count]
-            
-            # Combine frames side by side
-            combined_frame = np.hstack([
-                frame,
-                cv2.resize(ai_frame, (frame_width, frame_height))
-            ])
-            
-            # Display combined frame
-            frame_placeholder.image(
-                combined_frame,
-                channels="BGR",
-                use_column_width=True
-            )
-            
-            # Control playback speed
-            cv2.waitKey(int(1000 / (fps * play_speed)))
-            frame_count += 1
-            
+            # Add replay button
+            if frame_count >= len(frames):
+                if st.sidebar.button("Replay"):
+                    frame_count = 0
+        
         # Cleanup
         cap.release()
         os.unlink(tfile.name)