Upload 4 files

app.py

@@ -29,47 +29,78 @@ class AIDancePartner:
             ["Sync Partner", "Generate New Moves"]
         )
         
+        # Add playback controls
+        play_speed = st.sidebar.slider("Playback Speed", 0.1, 2.0, 1.0, 0.1)
+        
         if video_file:
-            self.process_video(video_file, mode)
+            self.process_video(video_file, mode, play_speed)
 
-    def process_video(self, video_file, mode):
-        # Create temporary file to store uploaded video
+    def process_video(self, video_file, mode, play_speed):
         tfile = tempfile.NamedTemporaryFile(delete=False)
         tfile.write(video_file.read())
         
-        # Process the video
         cap = cv2.VideoCapture(tfile.name)
         
-        # Display original and AI dance side by side
-        col1, col2 = st.columns(2)
+        # Get video properties
+        fps = cap.get(cv2.CAP_PROP_FPS)
+        frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
+        frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
+        total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
         
-        with col1:
-            st.header("Your Dance")
-            stframe1 = st.empty()
-            
-        with col2:
-            st.header("AI Partner")
-            stframe2 = st.empty()
-            
+        # 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
-                
-            # Detect pose in original frame
             pose_landmarks = self.pose_detector.detect_pose(frame)
+            all_poses.append(pose_landmarks)
+        
+        # Reset video capture
+        cap.set(cv2.CAP_PROP_POS_FRAMES, 0)
+        
+        # Generate AI dance sequence
+        ai_sequence = self.dance_generator.generate_dance_sequence(
+            all_poses, 
+            mode, 
+            total_frames,
+            (frame_height, frame_width)
+        )
+        
+        # Playback loop
+        frame_count = 0
+        while cap.isOpened():
+            ret, frame = cap.read()
+            if not ret:
+                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))
+            ])
             
-            if mode == "Sync Partner":
-                # Generate synchronized dance moves
-                ai_frame = self.dance_generator.sync_moves(pose_landmarks)
-            else:
-                # Generate new dance moves based on style
-                ai_frame = self.dance_generator.generate_new_moves(pose_landmarks)
+            # Display combined frame
+            frame_placeholder.image(
+                combined_frame,
+                channels="BGR",
+                use_column_width=True
+            )
             
-            # Visualize both frames
-            vis_frame = self.visualizer.draw_pose(frame, pose_landmarks)
-            stframe1.image(vis_frame, channels="BGR")
-            stframe2.image(ai_frame, channels="BGR")
+            # Control playback speed
+            cv2.waitKey(int(1000 / (fps * play_speed)))
+            frame_count += 1
             
         # Cleanup
         cap.release()

dance_generator.py

@@ -1,42 +1,185 @@
 import numpy as np
 import cv2
+from scipy.interpolate import interp1d
 
 class DanceGenerator:
     def __init__(self):
         self.prev_moves = []
         self.style_memory = []
+        self.avatar = cv2.imread('assets/dancer_avatar.png')  # Add a dancer avatar image
         
-    def sync_moves(self, pose_landmarks):
-        """Generate synchronized dance moves based on input pose"""
-        if pose_landmarks is None:
-            return np.zeros((480, 640, 3), dtype=np.uint8)
+    def generate_dance_sequence(self, all_poses, mode, total_frames, frame_size):
+        """Generate complete dance sequence for the entire video"""
+        height, width = frame_size
+        sequence = []
+        
+        if mode == "Sync Partner":
+            sequence = self._generate_sync_sequence(all_poses, total_frames, frame_size)
+        else:
+            sequence = self._generate_creative_sequence(all_poses, total_frames, frame_size)
             
-        # Convert landmarks to numpy array for processing
-        landmarks_array = self._landmarks_to_array(pose_landmarks)
+        return sequence
+        
+    def _generate_sync_sequence(self, all_poses, total_frames, frame_size):
+        """Generate synchronized dance sequence"""
+        height, width = frame_size
+        sequence = []
         
-        # Mirror the movements for sync mode
-        mirrored_moves = self._mirror_movements(landmarks_array)
+        # Convert all poses to arrays
+        pose_arrays = []
+        for pose in all_poses:
+            if pose is not None:
+                pose_arrays.append(self._landmarks_to_array(pose))
+            else:
+                pose_arrays.append(None)
+                
+        # Generate mirrored sequence with smooth transitions
+        for i in range(total_frames):
+            frame = np.zeros((height, width, 3), dtype=np.uint8)
+            
+            if pose_arrays[i] is not None:
+                # Mirror the pose
+                mirrored = self._mirror_movements(pose_arrays[i])
+                
+                # Add smooth transition from previous frame
+                if i > 0 and pose_arrays[i-1] is not None:
+                    mirrored = self._smooth_transition(pose_arrays[i-1], mirrored, 0.3)
+                
+                # Create dance frame
+                frame = self._create_enhanced_dance_frame(
+                    mirrored,
+                    frame_size,
+                    add_effects=True
+                )
+                
+            sequence.append(frame)
+            
+        return sequence
         
-        # Generate visualization frame
-        return self._create_dance_frame(mirrored_moves)
+    def _generate_creative_sequence(self, all_poses, total_frames, frame_size):
+        """Generate creative dance sequence based on style"""
+        height, width = frame_size
+        sequence = []
         
-    def generate_new_moves(self, pose_landmarks):
-        """Generate new dance moves based on learned style"""
-        if pose_landmarks is None:
-            return np.zeros((480, 640, 3), dtype=np.uint8)
+        # Analyze style from all poses
+        style_patterns = self._analyze_style_patterns(all_poses)
+        
+        # Generate new sequence using style patterns
+        for i in range(total_frames):
+            frame = np.zeros((height, width, 3), dtype=np.uint8)
+            
+            # Generate new pose based on style
+            new_pose = self._generate_style_based_pose(style_patterns, i/total_frames)
+            
+            if new_pose is not None:
+                frame = self._create_enhanced_dance_frame(
+                    new_pose,
+                    frame_size,
+                    add_effects=True
+                )
+                
+            sequence.append(frame)
             
-        # Convert landmarks to array
-        landmarks_array = self._landmarks_to_array(pose_landmarks)
+        return sequence
         
-        # Update style memory
-        self._update_style_memory(landmarks_array)
+    def _analyze_style_patterns(self, poses):
+        """Analyze dance style patterns from poses"""
+        patterns = []
         
-        # Generate new moves based on style
-        new_moves = self._generate_style_based_moves()
+        for pose in poses:
+            if pose is not None:
+                landmarks = self._landmarks_to_array(pose)
+                patterns.append(landmarks)
+                
+        return patterns
         
-        # Create visualization frame
-        return self._create_dance_frame(new_moves)
+    def _generate_style_based_pose(self, patterns, progress):
+        """Generate new pose based on style patterns and progress"""
+        if not patterns:
+            return None
+            
+        # Create smooth interpolation between poses
+        num_patterns = len(patterns)
+        pattern_idx = int(progress * (num_patterns - 1))
         
+        if pattern_idx < num_patterns - 1:
+            t = progress * (num_patterns - 1) - pattern_idx
+            pose = self._interpolate_poses(
+                patterns[pattern_idx],
+                patterns[pattern_idx + 1],
+                t
+            )
+        else:
+            pose = patterns[-1]
+            
+        return pose
+        
+    def _interpolate_poses(self, pose1, pose2, t):
+        """Smoothly interpolate between two poses"""
+        return pose1 * (1 - t) + pose2 * t
+        
+    def _create_enhanced_dance_frame(self, pose_array, frame_size, add_effects=True):
+        """Create enhanced visualization frame with effects"""
+        height, width = frame_size
+        frame = np.zeros((height, width, 3), dtype=np.uint8)
+        
+        # Convert coordinates
+        points = (pose_array[:, :2] * [width, height]).astype(int)
+        
+        # Draw enhanced skeleton
+        connections = self._get_pose_connections()
+        for connection in connections:
+            start_idx, end_idx = connection
+            if start_idx < len(points) and end_idx < len(points):
+                # Draw glowing lines
+                if add_effects:
+                    self._draw_glowing_line(
+                        frame,
+                        points[start_idx],
+                        points[end_idx],
+                        (0, 255, 0)
+                    )
+                else:
+                    cv2.line(frame, 
+                            tuple(points[start_idx]), 
+                            tuple(points[end_idx]),
+                            (0, 255, 0), 2)
+        
+        # Draw enhanced joints
+        for point in points:
+            if add_effects:
+                self._draw_glowing_point(frame, point, (0, 0, 255))
+            else:
+                cv2.circle(frame, tuple(point), 4, (0, 0, 255), -1)
+                
+        return frame
+        
+    def _draw_glowing_line(self, frame, start, end, color, thickness=2):
+        """Draw a line with glow effect"""
+        # Draw main line
+        cv2.line(frame, tuple(start), tuple(end), color, thickness)
+        
+        # Draw glow
+        for i in range(3):
+            alpha = 0.3 - i * 0.1
+            thickness = thickness + 2
+            cv2.line(frame, tuple(start), tuple(end),
+                    tuple([int(c * alpha) for c in color]),
+                    thickness)
+                    
+    def _draw_glowing_point(self, frame, point, color, radius=4):
+        """Draw a point with glow effect"""
+        # Draw main point
+        cv2.circle(frame, tuple(point), radius, color, -1)
+        
+        # Draw glow
+        for i in range(3):
+            alpha = 0.3 - i * 0.1
+            r = radius + i * 2
+            cv2.circle(frame, tuple(point), r,
+                      tuple([int(c * alpha) for c in color]),
+                      -1)
+
     def _landmarks_to_array(self, landmarks):
         """Convert MediaPipe landmarks to numpy array"""
         points = []