Update app.py

app.py

@@ -3,6 +3,9 @@ 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(
@@ -52,6 +55,18 @@ def local_css():
         </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()
     
@@ -70,6 +85,23 @@ def main():
         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'])
 
+        # Add video preview
+        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_column_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")
@@ -80,54 +112,61 @@ def main():
         """)
         st.markdown('</div>', unsafe_allow_html=True)
 
-    if uploaded_file is not None:
-        # Create a temporary file
-        with tempfile.NamedTemporaryFile(delete=False, suffix='.mp4') as tfile:
-            tfile.write(uploaded_file.read())
-            temp_input_path = tfile.name
-        
-        # Style selection with custom design
-        st.markdown("### Choose Your Dance Partner Style")
-        style = st.select_slider(
-            "",
-            options=["Sync Partner", "Creative Partner"],
-            value="Sync Partner"
-        )
-        
-        # Add description based on selected style
-        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:
-                # Create a progress bar
-                progress_bar = st.progress(0)
-                status_text = st.empty()
-                
-                # Processing steps
-                steps = ["Analyzing dance moves...", "Generating partner movements...", "Creating final video..."]
-                for i, step in enumerate(steps):
-                    status_text.text(step)
-                    progress_bar.progress((i + 1) * 33)
-                    time.sleep(1)
-                
-                # 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 in a nice container
-                st.markdown("### 🎥 Your Dance Duet")
-                st.video(output_path)
-                
-                # Download button with custom styling
-                col1, col2, col3 = st.columns([1, 2, 1])
-                with col2:
+        if uploaded_file is not None:
+            # Style selection with custom design
+            st.markdown("### 🎭 Choose Your Dance Partner Style")
+            style = st.select_slider(
+                "",
+                options=["Sync Partner", "Creative Partner"],
+                value="Sync Partner"
+            )
+            
+            # Add description based on selected style
+            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:
+                    # Create a progress bar
+                    progress_bar = st.progress(0)
+                    status_text = st.empty()
+                    
+                    # Processing steps with more detailed progress
+                    steps = [
+                        "Analyzing dance moves...",
+                        "Detecting pose landmarks...",
+                        "Generating partner movements...",
+                        "Applying style patterns...",
+                        "Creating final video..."
+                    ]
+                    
+                    for i, step in enumerate(steps):
+                        status_text.text(step)
+                        progress_bar.progress((i + 1) * 20)
+                        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 in a nice container
+                    st.markdown("### 🎥 Your Dance Duet")
+                    
+                    # Show preview of the output
+                    preview_output = get_video_preview(output_path)
+                    if preview_output:
+                        st.image(preview_output, use_column_width=True, caption="Dance Duet Preview")
+                    
+                    # Display the video
+                    st.video(output_path)
+                    
+                    # Download button with custom styling
                     with open(output_path, 'rb') as file:
                         st.download_button(
                             label="Download Your Dance Duet 📥",
@@ -135,21 +174,27 @@ def main():
                             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):
+                    
+                    # 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)
 
-    # Add footer
+    # Add footer with additional information
     st.markdown("""
         ---
-        <div style="text-align: center; color: #666;">
-            Made with ❤️ by AI Dance Partner Team
+        <div style="text-align: center;">
+            <h3>🌟 Features</h3>
+            <p>• Real-time pose detection</p>
+            <p>• Synchronized movement matching</p>
+            <p>• Creative dance style generation</p>
+            <p>• High-quality video output</p>
+            <br>
+            <p style="color: #666;">Made with ❤️ by AI Dance Partner Team</p>
         </div>
     """, unsafe_allow_html=True)