Upload 4 files

app.py

@@ -0,0 +1,83 @@
+import streamlit as st
+import cv2
+import numpy as np
+import tempfile
+import os
+from pose_detector import PoseDetector
+from dance_generator import DanceGenerator
+from dance_visualizer import DanceVisualizer
+
+class AIDancePartner:
+    def __init__(self):
+        self.pose_detector = PoseDetector()
+        self.dance_generator = DanceGenerator()
+        self.visualizer = DanceVisualizer()
+        
+    def setup_ui(self):
+        st.title("AI Dance Partner 💃🤖")
+        st.sidebar.header("Controls")
+        
+        # Upload video section
+        video_file = st.sidebar.file_uploader(
+            "Upload your dance video", 
+            type=['mp4', 'avi', 'mov']
+        )
+        
+        # Mode selection
+        mode = st.sidebar.selectbox(
+            "Select Mode",
+            ["Sync Partner", "Generate New Moves"]
+        )
+        
+        if video_file:
+            self.process_video(video_file, mode)
+
+    def process_video(self, video_file, mode):
+        # Create temporary file to store uploaded video
+        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)
+        
+        with col1:
+            st.header("Your Dance")
+            stframe1 = st.empty()
+            
+        with col2:
+            st.header("AI Partner")
+            stframe2 = st.empty()
+            
+        while cap.isOpened():
+            ret, frame = cap.read()
+            if not ret:
+                break
+                
+            # Detect pose in original frame
+            pose_landmarks = self.pose_detector.detect_pose(frame)
+            
+            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)
+            
+            # Visualize both frames
+            vis_frame = self.visualizer.draw_pose(frame, pose_landmarks)
+            stframe1.image(vis_frame, channels="BGR")
+            stframe2.image(ai_frame, channels="BGR")
+            
+        # Cleanup
+        cap.release()
+        os.unlink(tfile.name)
+
+def main():
+    app = AIDancePartner()
+    app.setup_ui()
+
+if __name__ == "__main__":
+    main()

dance_generator.py

@@ -0,0 +1,109 @@
+import numpy as np
+import cv2
+
+class DanceGenerator:
+    def __init__(self):
+        self.prev_moves = []
+        self.style_memory = []
+        
+    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)
+            
+        # Convert landmarks to numpy array for processing
+        landmarks_array = self._landmarks_to_array(pose_landmarks)
+        
+        # Mirror the movements for sync mode
+        mirrored_moves = self._mirror_movements(landmarks_array)
+        
+        # Generate visualization frame
+        return self._create_dance_frame(mirrored_moves)
+        
+    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)
+            
+        # Convert landmarks to array
+        landmarks_array = self._landmarks_to_array(pose_landmarks)
+        
+        # Update style memory
+        self._update_style_memory(landmarks_array)
+        
+        # Generate new moves based on style
+        new_moves = self._generate_style_based_moves()
+        
+        # Create visualization frame
+        return self._create_dance_frame(new_moves)
+        
+    def _landmarks_to_array(self, landmarks):
+        """Convert MediaPipe landmarks to numpy array"""
+        points = []
+        for landmark in landmarks.landmark:
+            points.append([landmark.x, landmark.y, landmark.z])
+        return np.array(points)
+        
+    def _mirror_movements(self, landmarks):
+        """Mirror the input movements"""
+        mirrored = landmarks.copy()
+        mirrored[:, 0] = 1 - mirrored[:, 0]  # Flip x coordinates
+        return mirrored
+        
+    def _update_style_memory(self, landmarks):
+        """Update memory of dance style"""
+        self.style_memory.append(landmarks)
+        if len(self.style_memory) > 30:  # Keep last 30 frames
+            self.style_memory.pop(0)
+            
+    def _generate_style_based_moves(self):
+        """Generate new moves based on learned style"""
+        if not self.style_memory:
+            return np.zeros((33, 3))  # Default pose shape
+            
+        # Simple implementation: interpolate between stored poses
+        base_pose = self.style_memory[-1]
+        if len(self.style_memory) > 1:
+            prev_pose = self.style_memory[-2]
+            t = np.random.random()
+            new_pose = t * base_pose + (1-t) * prev_pose
+        else:
+            new_pose = base_pose
+            
+        return new_pose
+        
+    def _create_dance_frame(self, pose_array):
+        """Create visualization frame from pose array"""
+        frame = np.zeros((480, 640, 3), dtype=np.uint8)
+        
+        # Convert normalized coordinates to pixel coordinates
+        points = (pose_array[:, :2] * [640, 480]).astype(int)
+        
+        # Draw connections between joints
+        connections = self._get_pose_connections()
+        for connection in connections:
+            start_idx, end_idx = connection
+            if start_idx < len(points) and end_idx < len(points):
+                cv2.line(frame, 
+                        tuple(points[start_idx]), 
+                        tuple(points[end_idx]),
+                        (0, 255, 0), 2)
+                
+        # Draw joints
+        for point in points:
+            cv2.circle(frame, tuple(point), 4, (0, 0, 255), -1)
+            
+        return frame
+        
+    def _get_pose_connections(self):
+        """Define connections between pose landmarks"""
+        return [
+            (0, 1), (1, 2), (2, 3), (3, 7),  # Face
+            (0, 4), (4, 5), (5, 6), (6, 8),
+            (9, 10), (11, 12), (11, 13), (13, 15),  # Arms
+            (12, 14), (14, 16),
+            (11, 23), (12, 24),  # Torso
+            (23, 24), (23, 25), (24, 26),  # Legs
+            (25, 27), (26, 28), (27, 29), (28, 30),
+            (29, 31), (30, 32)
+        ]

dance_visualizer.py

@@ -0,0 +1,27 @@
+import cv2
+import mediapipe as mp
+import numpy as np
+
+class DanceVisualizer:
+    def __init__(self):
+        self.mp_drawing = mp.solutions.drawing_utils
+        self.mp_pose = mp.solutions.pose
+        
+    def draw_pose(self, frame, pose_landmarks):
+        """Draw pose landmarks on frame"""
+        if pose_landmarks is None:
+            return frame
+            
+        # Create copy of frame
+        annotated_frame = frame.copy()
+        
+        # Draw the pose landmarks
+        self.mp_drawing.draw_landmarks(
+            annotated_frame,
+            pose_landmarks,
+            self.mp_pose.POSE_CONNECTIONS,
+            self.mp_drawing.DrawingSpec(color=(245,117,66), thickness=2, circle_radius=2),
+            self.mp_drawing.DrawingSpec(color=(245,66,230), thickness=2, circle_radius=2)
+        )
+        
+        return annotated_frame

pose_detector.py

@@ -0,0 +1,22 @@
+import mediapipe as mp
+import cv2
+import numpy as np
+
+class PoseDetector:
+    def __init__(self):
+        self.mp_pose = mp.solutions.pose
+        self.pose = self.mp_pose.Pose(
+            min_detection_confidence=0.5,
+            min_tracking_confidence=0.5
+        )
+        
+    def detect_pose(self, frame):
+        # Convert BGR to RGB
+        rgb_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+        
+        # Process the frame and detect poses
+        results = self.pose.process(rgb_frame)
+        
+        if results.pose_landmarks:
+            return results.pose_landmarks
+        return None