Update app.py

app.py

@@ -1,7 +1,8 @@
 """
-🍹 Video Background Replacer - IFRAME OPTIMIZED VERSION
+🍹 Video Background Replacer - IFRAME OPTIMIZED VERSION WITH 2-STAGE PIPELINE
 Optimized for embedding in MyAvatar Railway app with Claude's recommendations
 FIXED: Corrected MatAnyone API usage (processor.step() instead of processor.infer())
+RESTORED: 2-Stage Pipeline - MatAnyone → Green Screen → Custom Background
 """
 
 import streamlit as st
@@ -102,6 +103,16 @@
         margin: 20px 0;
     }
     
+    .stage-indicator {
+        background: linear-gradient(45deg, #FF6B6B, #4ECDC4);
+        color: white;
+        padding: 10px;
+        border-radius: 8px;
+        margin: 10px 0;
+        text-align: center;
+        font-weight: bold;
+    }
+    
     /* Iframe-specific optimizations */
     .block-container {
         padding-top: 1rem;
@@ -240,10 +251,10 @@ def process_frame_matanyone(self, frame, is_first_frame=False, mask=None):
             print(f"MatAnyone processing failed: {e}, using fallback")
             return self.create_simple_mask(frame).astype(np.float32) / 255.0
 
-    def process_frame(self, frame, background_image, is_first_frame=False, mask=None):
-        """Process a single frame with background replacement"""
+    def stage1_create_green_screen(self, frame, is_first_frame=False, mask=None):
+        """STAGE 1: Create green screen video using AI segmentation"""
         if self.use_matanyone:
-            # Use CORRECTED MatAnyone API
+            # Use MatAnyone for person segmentation
             alpha_mask = self.process_frame_matanyone(frame, is_first_frame, mask)
         elif self.use_mediapipe:
             try:
@@ -267,18 +278,66 @@ def process_frame(self, frame, background_image, is_first_frame=False, mask=None
             # Simple fallback method
             alpha_mask = self.create_simple_mask(frame).astype(np.float32) / 255.0
         
-        # Resize background to match frame
-        bg_resized = cv2.resize(background_image, (frame.shape[1], frame.shape[0]))
+        # Create pure green screen background
+        green_screen = np.full_like(frame, [0, 255, 0], dtype=np.uint8)  # Pure green (BGR)
+        
+        # Composite person on green screen
+        # alpha_mask values close to 1 = keep original (person)
+        # alpha_mask values close to 0 = use green screen (background)
+        green_screen_frame = frame * alpha_mask + green_screen * (1 - alpha_mask)
+        
+        return green_screen_frame.astype(np.uint8)
+
+    def stage2_chroma_key_replace(self, green_screen_frame, background_image):
+        """STAGE 2: Replace green screen with custom background using chroma key"""
+        try:
+            # Convert to HSV for better green screen detection
+            hsv = cv2.cvtColor(green_screen_frame, cv2.COLOR_BGR2HSV)
+            
+            # Define green screen range (optimized for pure green)
+            lower_green = np.array([35, 40, 40])   # Lower bound for green
+            upper_green = np.array([85, 255, 255]) # Upper bound for green
+            
+            # Create mask for green screen areas
+            green_mask = cv2.inRange(hsv, lower_green, upper_green)
+            
+            # Smooth the mask to avoid harsh edges
+            kernel = np.ones((3,3), np.uint8)
+            green_mask = cv2.morphologyEx(green_mask, cv2.MORPH_CLOSE, kernel)
+            green_mask = cv2.morphologyEx(green_mask, cv2.MORPH_OPEN, kernel)
+            
+            # Apply Gaussian blur for smoother edges
+            green_mask = cv2.GaussianBlur(green_mask, (5, 5), 0)
+            
+            # Convert mask to 3-channel and normalize
+            green_mask_3ch = cv2.cvtColor(green_mask, cv2.COLOR_GRAY2BGR).astype(np.float32) / 255.0
+            
+            # Resize background to match frame
+            bg_resized = cv2.resize(background_image, (green_screen_frame.shape[1], green_screen_frame.shape[0]))
+            
+            # Replace green areas with background
+            # green_mask_3ch = 1 where green (use background), 0 where not green (keep original)
+            result = green_screen_frame * (1 - green_mask_3ch) + bg_resized * green_mask_3ch
+            
+            return result.astype(np.uint8)
+            
+        except Exception as e:
+            print(f"Stage 2 chroma key failed: {e}")
+            # Fallback: return original frame
+            return green_screen_frame
+
+    def process_frame_2stage(self, frame, background_image, is_first_frame=False, mask=None):
+        """Process a single frame with 2-stage background replacement"""
+        # STAGE 1: Create green screen
+        green_screen_frame = self.stage1_create_green_screen(frame, is_first_frame, mask)
         
-        # Apply background replacement
-        # mask values close to 1 = keep original (person)
-        # mask values close to 0 = use background
-        result = frame * alpha_mask + bg_resized * (1 - alpha_mask)
+        # STAGE 2: Replace green screen with custom background
+        final_frame = self.stage2_chroma_key_replace(green_screen_frame, background_image)
         
-        return result.astype(np.uint8)
+        return final_frame
 
     def process_video(self, video_path, background_image, progress_callback=None):
-        """Process entire video with background replacement"""
+        """Process entire video with 2-stage background replacement"""
         # Convert PIL background to OpenCV format
         if isinstance(background_image, Image.Image):
             background_image = cv2.cvtColor(np.array(background_image), cv2.COLOR_RGB2BGR)
@@ -309,8 +368,8 @@ def process_video(self, video_path, background_image, progress_callback=None):
                 is_first_frame = (frame_count == 0)
                 mask = None  # Could add automatic mask detection here if needed
                 
-                # Process frame
-                processed_frame = self.process_frame(frame, background_image, is_first_frame, mask)
+                # Process frame with 2-stage pipeline
+                processed_frame = self.process_frame_2stage(frame, background_image, is_first_frame, mask)
                 out.write(processed_frame)
                 
                 frame_count += 1
@@ -318,7 +377,8 @@ def process_video(self, video_path, background_image, progress_callback=None):
                 # Update progress
                 if progress_callback:
                     progress = frame_count / total_frames
-                    progress_callback(progress, frame_count, total_frames)
+                    stage = "stage1" if progress < 0.5 else "stage2"
+                    progress_callback(stage, progress, frame_count, total_frames)
         
         finally:
             cap.release()
@@ -329,20 +389,30 @@ def process_video(self, video_path, background_image, progress_callback=None):
 def main():
     # Compact header for iframe
     st.markdown('<h1 class="main-header">🍹 Video Background Replacer</h1>', unsafe_allow_html=True)
-    st.markdown('<p style="text-align: center; font-size: 1.1rem; color: #666; margin-bottom: 2rem;">Replace your video background with AI!</p>', unsafe_allow_html=True)
+    st.markdown('<p style="text-align: center; font-size: 1.1rem; color: #666; margin-bottom: 2rem;">2-Stage AI Processing: Person Segmentation → Green Screen → Custom Background</p>', unsafe_allow_html=True)
     
     # Initialize the background replacer with loading feedback
     if 'replacer' not in st.session_state:
-        with st.spinner('🔄 Loading AI Processing Interface...'):
+        with st.spinner('🔄 Loading 2-Stage AI Processing Interface...'):
             st.session_state.replacer = BackgroundReplacer()
     
     # Show initialization status
     if hasattr(st.session_state.replacer, 'use_matanyone') and st.session_state.replacer.use_matanyone:
-        st.success('🚀 MatAnyone AI Ready - Ultimate Quality!')
+        st.success('🚀 MatAnyone AI Ready - 2-Stage Ultimate Quality!')
     elif hasattr(st.session_state.replacer, 'use_mediapipe') and st.session_state.replacer.use_mediapipe:
-        st.success('🎯 MediaPipe AI Ready - Good Quality!')
+        st.success('🎯 MediaPipe AI Ready - 2-Stage Good Quality!')
     else:
-        st.info('📱 Basic Processing Ready')
+        st.info('📱 Basic 2-Stage Processing Ready')
+    
+    # Show 2-stage process explanation
+    st.markdown("""
+    <div class="processing-box">
+        <h4>🔄 2-Stage Processing Pipeline:</h4>
+        <p><strong>Stage 1:</strong> AI extracts person from video → Creates clean green screen video</p>
+        <p><strong>Stage 2:</strong> Chroma key replacement → Green screen becomes your custom background</p>
+        <p><em>This approach ensures the highest quality and most reliable results!</em></p>
+    </div>
+    """, unsafe_allow_html=True)
     
     # Compact upload section
     col1, col2 = st.columns(2)
@@ -374,10 +444,10 @@ def main():
     # Processing section - SIMPLIFIED to avoid DOM manipulation issues
     if uploaded_video and uploaded_background:
         st.markdown("---")
-        st.markdown("### 🚀 Ready to Process!")
+        st.markdown("### 🚀 Ready for 2-Stage Processing!")
         
         # Process button (clean and stable)
-        if st.button("🍹 PROCESS VIDEO", key="process_button", use_container_width=True):
+        if st.button("🍹 START 2-STAGE PROCESSING", key="process_button", use_container_width=True):
             # Save uploaded files
             with tempfile.NamedTemporaryFile(delete=False, suffix='.mp4') as tmp_video:
                 tmp_video.write(uploaded_video.read())
@@ -390,30 +460,37 @@ def main():
             
             with result_container:
                 st.markdown('<div class="processing-box">', unsafe_allow_html=True)
-                st.markdown("### 🔄 Processing Your Video...")
+                st.markdown("### 🔄 2-Stage Processing Your Video...")
                 
                 # Progress tracking
                 progress_bar = st.progress(0)
                 status_text = st.empty()
+                stage_indicator = st.empty()
                 
-                def update_progress(progress, frame_count, total_frames):
+                def update_progress(stage, progress, frame_count, total_frames):
                     progress_bar.progress(progress)
                     status_text.text(f"Processing frame {frame_count}/{total_frames} ({progress*100:.1f}%)")
+                    
+                    if stage == "stage1":
+                        stage_indicator.markdown('<div class="stage-indicator">🎬 STAGE 1: Creating Green Screen Video</div>', unsafe_allow_html=True)
+                    elif stage == "stage2":
+                        stage_indicator.markdown('<div class="stage-indicator">🖼️ STAGE 2: Applying Custom Background</div>', unsafe_allow_html=True)
                 
                 try:
                     # Process the video
-                    st.info("🚀 Starting video processing...")
+                    st.info("🚀 Starting 2-stage video processing...")
                     output_path = st.session_state.replacer.process_video(
                         video_path, background_image, update_progress
                     )
                     
                     # Complete progress
                     progress_bar.progress(1.0)
-                    status_text.text("✅ Processing complete!")
+                    status_text.text("✅ 2-Stage Processing complete!")
+                    stage_indicator.markdown('<div class="stage-indicator">🎉 BOTH STAGES COMPLETED!</div>', unsafe_allow_html=True)
                     st.markdown('</div>', unsafe_allow_html=True)
                     
                     # Success message
-                    st.markdown('<div class="success-box">🎉 Video Successfully Processed! 🎉</div>', unsafe_allow_html=True)
+                    st.markdown('<div class="success-box">🎉 Video Successfully Processed with 2-Stage Pipeline! 🎉</div>', unsafe_allow_html=True)
                     
                     # Display result immediately in same container
                     if os.path.exists(output_path) and os.path.getsize(output_path) > 0:
@@ -429,13 +506,13 @@ def update_progress(progress, frame_count, total_frames):
                         st.download_button(
                             label="⬇️ Download Processed Video",
                             data=video_data,
-                            file_name=f"background_replaced_{int(time.time())}.mp4",
+                            file_name=f"2stage_background_replaced_{int(time.time())}.mp4",
                             mime="video/mp4",
                             use_container_width=True,
                             key=f"download_button_{int(time.time())}"
                         )
                         
-                        st.success("✅ Video ready for download!")
+                        st.success("✅ 2-Stage processed video ready for download!")
                         
                         # Cleanup temp files
                         try:
@@ -448,18 +525,18 @@ def update_progress(progress, frame_count, total_frames):
                         st.error("❌ Output video file is empty or corrupted")
                         
                 except Exception as e:
-                    status_text.text("❌ Processing failed")
+                    status_text.text("❌ 2-Stage processing failed")
                     st.error(f"❌ Processing failed: {str(e)}")
                     st.info("💡 Try with a shorter video or different background image")
     
     else:
-        st.info("👆 Upload both a video and background image to start processing!")
+        st.info("👆 Upload both a video and background image to start 2-stage processing!")
     
     # Compact footer
     st.markdown("---")
     st.markdown("""
     <div style="text-align: center; color: #666; padding: 10px;">
-        <p><small>🍹 Powered by MatAnyone and MediaPipe | Optimized for MyAvatar | API Fixed</small></p>
+        <p><small>🍹 2-Stage Pipeline: MatAnyone → Green Screen → Custom Background | Optimized for MyAvatar</small></p>
     </div>
     """, unsafe_allow_html=True)