app.py

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()