File size: 6,080 Bytes
ec27c1f 96fe8c7 |
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 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 |
import streamlit as st
from colab import AIDancePartner
import tempfile
import os
import time
import cv2
from PIL import Image
import io
# Set page configuration
st.set_page_config(
page_title="AI Dance Partner",
page_icon="💃",
layout="wide",
initial_sidebar_state="expanded"
)
# Custom CSS for better styling
def local_css():
st.markdown("""
<style>
.main {
padding: 2rem;
}
.stButton>button {
background-color: #FF4B4B;
color: white;
border-radius: 20px;
padding: 0.5rem 2rem;
font-weight: bold;
}
.stButton>button:hover {
background-color: #FF6B6B;
border-color: #FF4B4B;
}
.upload-text {
font-size: 1.2rem;
color: #666;
margin-bottom: 1rem;
}
.title-container {
background: linear-gradient(90deg, #FF4B4B, #FF8C8C);
padding: 2rem;
border-radius: 10px;
margin-bottom: 2rem;
color: white;
text-align: center;
}
.info-box {
background-color: #f0f2f6;
padding: 1rem;
border-radius: 10px;
margin-bottom: 1rem;
}
</style>
""", unsafe_allow_html=True)
def get_video_preview(video_path):
"""Generate a preview frame from the video"""
cap = cv2.VideoCapture(video_path)
ret, frame = cap.read()
cap.release()
if ret:
# Convert BGR to RGB
frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
return Image.fromarray(frame)
return None
def main():
local_css()
# Title section with gradient background
st.markdown("""
<div class="title-container">
<h1>🕺 AI Dance Partner 💃</h1>
<p style="font-size: 1.2rem;">Transform your solo dance into a dynamic duet!</p>
</div>
""", unsafe_allow_html=True)
# Create two columns for layout
col1, col2 = st.columns([2, 1])
with col1:
st.markdown('<p class="upload-text">Upload your dance video and watch the magic happen!</p>', unsafe_allow_html=True)
uploaded_file = st.file_uploader("", type=['mp4', 'avi', 'mov'])
if uploaded_file is not None:
# Create a temporary file for the uploaded video
with tempfile.NamedTemporaryFile(delete=False, suffix='.mp4') as tfile:
tfile.write(uploaded_file.read())
temp_input_path = tfile.name
# Show video preview
st.markdown("### 📽️ Preview")
preview_image = get_video_preview(temp_input_path)
if preview_image:
st.image(preview_image, use_container_width=True, caption="Video Preview")
# Add video player for original
st.markdown("### 🎥 Original Video")
st.video(temp_input_path)
with col2:
st.markdown('<div class="info-box">', unsafe_allow_html=True)
st.markdown("### How it works")
st.markdown("""
1. Upload your solo dance video
2. Choose your preferred dance style
3. Watch as AI creates your perfect dance partner!
""")
st.markdown('</div>', unsafe_allow_html=True)
if uploaded_file is not None:
st.markdown("### 🎭 Choose Your Dance Partner Style")
style = st.select_slider(
"",
options=["Sync Partner", "Creative Partner"],
value="Sync Partner"
)
if style == "Sync Partner":
st.info("💫 Sync Partner will mirror your movements in perfect harmony.")
else:
st.info("🎨 Creative Partner will add its own artistic flair to your dance.")
if st.button("Generate Dance Partner 🎬"):
try:
progress_bar = st.progress(0)
status_text = st.empty()
steps = [
"Analyzing dance moves...",
"Detecting pose landmarks...",
"Generating partner movements...",
"Creating final video..."
]
for i, step in enumerate(steps):
status_text.text(step)
progress_bar.progress((i + 1) * 25)
time.sleep(0.5)
# Process video
dance_partner = AIDancePartner()
output_path = dance_partner.process_video(temp_input_path, mode=style)
# Update progress
progress_bar.progress(100)
status_text.text("Done! 🎉")
# Display result
st.markdown("### 🎥 Your Dance Duet")
st.video(output_path)
# Download button
with open(output_path, 'rb') as file:
st.download_button(
label="Download Video 📥",
data=file,
file_name="ai_dance_partner.mp4",
mime="video/mp4"
)
# Cleanup temporary files
os.unlink(temp_input_path)
os.unlink(output_path)
except Exception as e:
st.error(f"Oops! Something went wrong: {str(e)}")
if os.path.exists(temp_input_path):
os.unlink(temp_input_path)
if __name__ == "__main__":
main() |