Upload colab.py

colab.py

@@ -0,0 +1,469 @@
+
+# Import necessary libraries
+import cv2
+import mediapipe as mp
+import numpy as np
+from scipy.interpolate import interp1d
+from google.colab import files
+from IPython.display import HTML, display
+from IPython.display import clear_output
+import time
+import base64
+
+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):
+        rgb_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+        results = self.pose.process(rgb_frame)
+        return results.pose_landmarks if results.pose_landmarks else None
+
+class DanceGenerator:
+    def __init__(self):
+        self.prev_moves = []
+        self.style_memory = []
+        self.rhythm_patterns = []
+
+    def generate_dance_sequence(self, all_poses, mode, total_frames, frame_size):
+        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)
+
+        return sequence
+
+    def _generate_sync_sequence(self, all_poses, total_frames, frame_size):
+        height, width = frame_size
+        sequence = []
+        
+        # Enhanced rhythm analysis
+        rhythm_window = 10  # Analyze chunks of frames for rhythm
+        beat_positions = self._detect_dance_beats(all_poses, rhythm_window)
+        
+        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)
+                
+        for i in range(total_frames):
+            frame = np.zeros((height, width, 3), dtype=np.uint8)
+            
+            if pose_arrays[i] is not None:
+                # Enhanced mirroring with rhythm awareness
+                mirrored = self._mirror_movements(pose_arrays[i])
+                
+                # Apply rhythm-based movement enhancement
+                if i in beat_positions:
+                    mirrored = self._enhance_movement_on_beat(mirrored)
+                
+                if i > 0 and pose_arrays[i-1] is not None:
+                    mirrored = self._smooth_transition(pose_arrays[i-1], mirrored, 0.3)
+                
+                frame = self._create_enhanced_dance_frame(
+                    mirrored,
+                    frame_size,
+                    add_effects=True
+                )
+            
+            sequence.append(frame)
+        
+        return sequence
+        
+    def _detect_dance_beats(self, poses, window_size):
+        """Detect main beats in the dance sequence"""
+        beat_positions = []
+        
+        if len(poses) < window_size:
+            return beat_positions
+            
+        for i in range(window_size, len(poses)):
+            if poses[i] is not None and poses[i-1] is not None:
+                curr_pose = self._landmarks_to_array(poses[i])
+                prev_pose = self._landmarks_to_array(poses[i-1])
+                
+                # Calculate movement magnitude
+                movement = np.mean(np.abs(curr_pose - prev_pose))
+                
+                # Detect significant movements as beats
+                if movement > np.mean(self.rhythm_patterns) + np.std(self.rhythm_patterns):
+                    beat_positions.append(i)
+                    
+        return beat_positions
+        
+    def _enhance_movement_on_beat(self, pose):
+        """Enhance movements during detected beats"""
+        # Amplify movements slightly on beats
+        center = np.mean(pose, axis=0)
+        enhanced_pose = pose.copy()
+        
+        for i in range(len(pose)):
+            # Amplify movement relative to center
+            vector = pose[i] - center
+            enhanced_pose[i] = center + vector * 1.2
+            
+        return enhanced_pose
+
+    def _generate_creative_sequence(self, all_poses, total_frames, frame_size):
+        """Generate creative dance sequence based on style"""
+        height, width = frame_size
+        sequence = []
+
+        # 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)
+
+        return sequence
+
+    def _analyze_style_patterns(self, poses):
+        """Enhanced style analysis including rhythm and movement patterns"""
+        patterns = []
+        rhythm_data = []
+        
+        for i in range(1, len(poses)):
+            if poses[i] is not None and poses[i-1] is not None:
+                # Calculate movement speed and direction
+                curr_pose = self._landmarks_to_array(poses[i])
+                prev_pose = self._landmarks_to_array(poses[i-1])
+                
+                # Analyze movement velocity
+                velocity = np.mean(np.abs(curr_pose - prev_pose), axis=0)
+                rhythm_data.append(velocity)
+                
+                # Store enhanced pattern data
+                pattern_info = {
+                    'pose': curr_pose,
+                    'velocity': velocity,
+                    'acceleration': velocity if i == 1 else velocity - prev_velocity
+                }
+                patterns.append(pattern_info)
+                prev_velocity = velocity
+                
+        self.rhythm_patterns = rhythm_data
+        return patterns
+
+    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 = []
+        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)
+        ]
+
+    def _smooth_transition(self, prev_pose, current_pose, smoothing_factor=0.3):
+        """Create smooth transition between poses"""
+        if prev_pose is None or current_pose is None:
+            return current_pose
+
+        # Interpolate between previous and current pose
+        smoothed_pose = (1 - smoothing_factor) * prev_pose + smoothing_factor * current_pose
+
+        # Ensure the smoothed pose maintains proper proportions
+        # Normalize joint positions relative to hip center
+        hip_center_idx = 23  # Index for hip center landmark
+
+        prev_hip = prev_pose[hip_center_idx]
+        current_hip = current_pose[hip_center_idx]
+        smoothed_hip = smoothed_pose[hip_center_idx]
+
+        # Adjust positions relative to hip center
+        for i in range(len(smoothed_pose)):
+            if i != hip_center_idx:
+                # Calculate relative positions
+                prev_relative = prev_pose[i] - prev_hip
+                current_relative = current_pose[i] - current_hip
+
+                # Interpolate relative positions
+                smoothed_relative = (1 - smoothing_factor) * prev_relative + smoothing_factor * current_relative
+
+                # Update smoothed pose
+                smoothed_pose[i] = smoothed_hip + smoothed_relative
+
+        return smoothed_pose
+
+class AIDancePartner:
+    def __init__(self):
+        self.pose_detector = PoseDetector()
+        self.dance_generator = DanceGenerator()
+
+    def process_video(self, video_path, mode="Sync Partner"):
+        cap = cv2.VideoCapture(video_path)
+
+        # Get video properties
+        fps = int(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))
+
+        # Create output video writer
+        output_path = 'output_dance.mp4'
+        fourcc = cv2.VideoWriter_fourcc(*'mp4v')
+        out = cv2.VideoWriter(output_path, fourcc, fps,
+                            (frame_width * 2, frame_height))
+
+        # Pre-process video to extract all poses
+        print("Processing video...")
+        all_poses = []
+        frame_count = 0
+
+        while cap.isOpened():
+            ret, frame = cap.read()
+            if not ret:
+                break
+
+            pose_landmarks = self.pose_detector.detect_pose(frame)
+            all_poses.append(pose_landmarks)
+
+            frame_count += 1
+            if frame_count % 30 == 0:
+                print(f"Processed {frame_count}/{total_frames} frames")
+
+        # Generate AI dance sequence
+        print("Generating AI dance sequence...")
+        ai_sequence = self.dance_generator.generate_dance_sequence(
+            all_poses,
+            mode,
+            total_frames,
+            (frame_height, frame_width)
+        )
+
+        # Reset video capture and create final video
+        cap.set(cv2.CAP_PROP_POS_FRAMES, 0)
+        frame_count = 0
+        print("Creating final video...")
+
+        while cap.isOpened():
+            ret, frame = cap.read()
+            if not ret:
+                break
+
+            # Get corresponding AI frame
+            ai_frame = ai_sequence[frame_count]
+
+            # Combine frames side by side
+            combined_frame = np.hstack([frame, ai_frame])
+
+            # Write frame to output video
+            out.write(combined_frame)
+
+            frame_count += 1
+            if frame_count % 30 == 0:
+                print(f"Processed {frame_count}/{total_frames} frames")
+
+        # Release resources
+        cap.release()
+        out.release()
+
+        print("Video processing complete!")
+        return output_path
+
+def upload_and_process_video():
+    """Function to handle video upload and processing in Colab"""
+    print("Please upload a video file...")
+    uploaded = files.upload()
+
+    if uploaded:
+        video_path = list(uploaded.keys())[0]
+        print(f"Processing video: {video_path}")
+
+        # Create AI Dance Partner instance
+        dance_partner = AIDancePartner()
+
+        # Process video
+        output_path = dance_partner.process_video(video_path)
+
+        # Display the output video
+        def create_video_player(video_path):
+            video_html = f'''
+                <video width="100%" height="480" controls>
+                    <source src="data:video/mp4;base64,{base64.b64encode(open(output_path, 'rb').read()).decode()}" type="video/mp4">
+                    Your browser does not support the video tag.
+                </video>
+            '''
+            return HTML(video_html)
+
+        print("Displaying processed video...")
+        display(create_video_player(output_path))
+
+        # Option to download the processed video
+        files.download(output_path)
+
+# Run the application
+if __name__ == "__main__":
+    print("AI Dance Partner - Video Processing")
+    print("==================================")
+    upload_and_process_video()