File size: 2,582 Bytes
75aee9e |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 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() |