Update app.py

app.py

@@ -1,6 +1,7 @@
 """
 🍹 Video Background Replacer - IFRAME OPTIMIZED VERSION
 Optimized for embedding in MyAvatar Railway app with Claude's recommendations
+FIXED: Corrected MatAnyone API usage (processor.step() instead of processor.infer())
 """
 
 import streamlit as st
@@ -124,43 +125,50 @@ def __init__(self):
         """Initialize the background replacer with MatAnyone or MediaPipe"""
         self.use_matanyone = False
         self.use_mediapipe = False
+        self.matanyone_processor = None
+        self.matanyone_model = None
+        self.device = None
         
         print("🔄 Starting AI model initialization...")
         
-        # Try MatAnyone first (best quality)
+        # Try MatAnyone first (best quality) with CORRECTED API
         try:
             print("🚀 Attempting to load MatAnyone...")
-            from matanyone import InferenceCore
+            from matanyone.inference.inference_core import InferenceCore
+            from matanyone.utils.get_default_model import get_matanyone_model
+            from matanyone.utils.device import get_default_device
             
             print("📦 MatAnyone imported successfully, initializing processor...")
-            # Initialize MatAnyone with model name (HuggingFace style)
-            self.matanyone_processor = InferenceCore("PeiqingYang/MatAnyone")
+            
+            # Get device
+            self.device = get_default_device()
+            
+            # Load MatAnyone model (auto-downloads if needed)
+            self.matanyone_model = get_matanyone_model(None, self.device)
+            
+            # Initialize processor with CORRECT API
+            self.matanyone_processor = InferenceCore(self.matanyone_model, cfg=self.matanyone_model.cfg)
+            
             self.use_matanyone = True
             print("✅ MatAnyone AI loaded successfully!")
             
         except ImportError as e:
             print(f"⚠️ MatAnyone not found ({str(e)}), trying MediaPipe...")
-            # Fallback to MediaPipe
-            try:
-                import mediapipe as mp
-                self.mp_selfie_segmentation = mp.solutions.selfie_segmentation
-                self.selfie_segmentation = self.mp_selfie_segmentation.SelfieSegmentation(model_selection=1)
-                self.use_mediapipe = True
-                print("✅ MediaPipe AI loaded successfully!")
-            except ImportError:
-                print("📱 Using basic background replacement (install MatAnyone or MediaPipe for better results)")
-        
+            self._init_mediapipe()
         except Exception as e:
             print(f"⚠️ MatAnyone failed to load ({str(e)}), trying MediaPipe...")
-            # Fallback to MediaPipe
-            try:
-                import mediapipe as mp
-                self.mp_selfie_segmentation = mp.solutions.selfie_segmentation
-                self.selfie_segmentation = self.mp_selfie_segmentation.SelfieSegmentation(model_selection=1)
-                self.use_mediapipe = True
-                print("✅ MediaPipe AI loaded successfully!")
-            except ImportError:
-                print("📱 Using basic background replacement")
+            self._init_mediapipe()
+    
+    def _init_mediapipe(self):
+        """Initialize MediaPipe as fallback"""
+        try:
+            import mediapipe as mp
+            self.mp_selfie_segmentation = mp.solutions.selfie_segmentation
+            self.selfie_segmentation = self.mp_selfie_segmentation.SelfieSegmentation(model_selection=1)
+            self.use_mediapipe = True
+            print("✅ MediaPipe AI loaded successfully!")
+        except ImportError:
+            print("📱 Using basic background replacement (install MatAnyone or MediaPipe for better results)")
     
     def create_simple_mask(self, frame):
         """Create a simple background mask using color detection"""
@@ -188,31 +196,55 @@ def create_simple_mask(self, frame):
         
         return mask
 
-    def process_frame(self, frame, background_image):
-        """Process a single frame with background replacement"""
-        if self.use_matanyone:
-            try:
-                # Convert frame to PIL Image for MatAnyone
-                frame_pil = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
-                
-                # Use MatAnyone for segmentation
-                result = self.matanyone_processor.infer(frame_pil)
-                
-                # Extract the mask from MatAnyone result
-                if hasattr(result, 'alpha') and result.alpha is not None:
-                    # MatAnyone returns alpha matte
-                    mask = np.array(result.alpha)
-                    if len(mask.shape) == 2:
-                        mask = np.stack([mask] * 3, axis=-1)
-                    mask = mask.astype(np.float32) / 255.0
-                else:
-                    # Fallback if alpha not available
-                    mask = self.create_simple_mask(frame).astype(np.float32) / 255.0
+    def process_frame_matanyone(self, frame, is_first_frame=False, mask=None):
+        """Process frame using CORRECTED MatAnyone API"""
+        try:
+            # Convert frame to tensor format expected by MatAnyone
+            # Frame should be in RGB format and normalized to [0,1]
+            if len(frame.shape) == 3 and frame.shape[2] == 3:
+                # Convert BGR to RGB
+                frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+            else:
+                frame_rgb = frame
+            
+            # Convert to tensor and normalize
+            frame_tensor = torch.from_numpy(frame_rgb).permute(2, 0, 1).float() / 255.0
+            frame_tensor = frame_tensor.to(self.device)
+            
+            if is_first_frame and mask is not None:
+                # First frame: provide mask
+                mask_tensor = torch.from_numpy(mask).float().to(self.device)
+                if len(mask_tensor.shape) == 3:
+                    mask_tensor = mask_tensor[:, :, 0]  # Take first channel if 3-channel
                 
-            except Exception as e:
-                print(f"MatAnyone processing failed: {e}, using fallback")
-                mask = self.create_simple_mask(frame).astype(np.float32) / 255.0
+                objects = [1]  # Object ID
                 
+                # CORRECTED API: Use .step() not .infer()
+                output_prob = self.matanyone_processor.step(frame_tensor, mask_tensor, objects=objects)
+                output_prob = self.matanyone_processor.step(frame_tensor, first_frame_pred=True)
+            else:
+                # Subsequent frames: no mask needed
+                output_prob = self.matanyone_processor.step(frame_tensor)
+            
+            # Convert output to mask
+            alpha_mask = self.matanyone_processor.output_prob_to_mask(output_prob)
+            alpha_mask = alpha_mask.cpu().numpy()
+            
+            # Ensure mask is 3-channel and proper format
+            if len(alpha_mask.shape) == 2:
+                alpha_mask = np.stack([alpha_mask] * 3, axis=-1)
+            
+            return alpha_mask.astype(np.float32)
+            
+        except Exception as e:
+            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"""
+        if self.use_matanyone:
+            # Use CORRECTED MatAnyone API
+            alpha_mask = self.process_frame_matanyone(frame, is_first_frame, mask)
         elif self.use_mediapipe:
             try:
                 # Convert BGR to RGB for MediaPipe
@@ -223,17 +255,17 @@ def process_frame(self, frame, background_image):
                 
                 if results.segmentation_mask is not None:
                     # Convert segmentation mask to 3-channel
-                    mask = results.segmentation_mask
-                    mask = np.stack([mask] * 3, axis=-1).astype(np.float32)
+                    alpha_mask = results.segmentation_mask
+                    alpha_mask = np.stack([alpha_mask] * 3, axis=-1).astype(np.float32)
                 else:
-                    mask = self.create_simple_mask(frame).astype(np.float32) / 255.0
+                    alpha_mask = self.create_simple_mask(frame).astype(np.float32) / 255.0
                     
             except Exception as e:
                 print(f"MediaPipe processing failed: {e}, using fallback")
-                mask = self.create_simple_mask(frame).astype(np.float32) / 255.0
+                alpha_mask = self.create_simple_mask(frame).astype(np.float32) / 255.0
         else:
             # Simple fallback method
-            mask = self.create_simple_mask(frame).astype(np.float32) / 255.0
+            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]))
@@ -241,7 +273,7 @@ def process_frame(self, frame, background_image):
         # Apply background replacement
         # mask values close to 1 = keep original (person)
         # mask values close to 0 = use background
-        result = frame * mask + bg_resized * (1 - mask)
+        result = frame * alpha_mask + bg_resized * (1 - alpha_mask)
         
         return result.astype(np.uint8)
 
@@ -273,8 +305,12 @@ def process_video(self, video_path, background_image, progress_callback=None):
                 if not ret:
                     break
                 
+                # For MatAnyone, first frame needs special handling
+                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)
+                processed_frame = self.process_frame(frame, background_image, is_first_frame, mask)
                 out.write(processed_frame)
                 
                 frame_count += 1
@@ -335,17 +371,13 @@ def main():
         if uploaded_background:
             st.image(uploaded_background, caption="Background Preview")
     
-    # Processing section
+    # Processing section - SIMPLIFIED to avoid DOM manipulation issues
     if uploaded_video and uploaded_background:
         st.markdown("---")
         st.markdown("### 🚀 Ready to Process!")
         
         # Process button (clean and stable)
         if st.button("🍹 PROCESS VIDEO", key="process_button", use_container_width=True):
-            # Clear any previous results first
-            if 'video_result' in st.session_state:
-                del st.session_state['video_result']
-            
             # Save uploaded files
             with tempfile.NamedTemporaryFile(delete=False, suffix='.mp4') as tmp_video:
                 tmp_video.write(uploaded_video.read())
@@ -353,10 +385,10 @@ def main():
             
             background_image = Image.open(uploaded_background)
             
-            # Processing container - clean and professional
-            processing_container = st.empty()
+            # Create stable container for processing and results
+            result_container = st.container()
             
-            with processing_container.container():
+            with result_container:
                 st.markdown('<div class="processing-box">', unsafe_allow_html=True)
                 st.markdown("### 🔄 Processing Your Video...")
                 
@@ -375,16 +407,35 @@ def update_progress(progress, frame_count, total_frames):
                         video_path, background_image, update_progress
                     )
                     
-                    # Clear processing animation
-                    processing_container.empty()
+                    # Complete progress
+                    progress_bar.progress(1.0)
+                    status_text.text("✅ Processing complete!")
+                    st.markdown('</div>', unsafe_allow_html=True)
                     
                     # Success message
                     st.markdown('<div class="success-box">🎉 Video Successfully Processed! 🎉</div>', unsafe_allow_html=True)
                     
-                    # Store result in session state
+                    # Display result immediately in same container
                     if os.path.exists(output_path) and os.path.getsize(output_path) > 0:
                         with open(output_path, 'rb') as video_file:
-                            st.session_state['video_result'] = video_file.read()
+                            video_data = video_file.read()
+                        
+                        st.markdown("### 🎬 Your Processed Video:")
+                        
+                        # Show the video
+                        st.video(video_data)
+                        
+                        # Download button
+                        st.download_button(
+                            label="⬇️ Download Processed Video",
+                            data=video_data,
+                            file_name=f"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!")
                         
                         # Cleanup temp files
                         try:
@@ -397,28 +448,9 @@ def update_progress(progress, frame_count, total_frames):
                         st.error("❌ Output video file is empty or corrupted")
                         
                 except Exception as e:
-                    processing_container.empty()
+                    status_text.text("❌ Processing failed")
                     st.error(f"❌ Processing failed: {str(e)}")
                     st.info("💡 Try with a shorter video or different background image")
-        
-        # Display results if they exist (separate from processing)
-        if 'video_result' in st.session_state:
-            st.markdown("### 🎬 Your Processed Video:")
-            
-            # Show the video
-            st.video(st.session_state['video_result'])
-            
-            # Download button
-            st.download_button(
-                label="⬇️ Download Processed Video",
-                data=st.session_state['video_result'],
-                file_name=f"background_replaced_{int(time.time())}.mp4",
-                mime="video/mp4",
-                use_container_width=True,
-                key="download_button"
-            )
-            
-            st.success("✅ Video ready for download!")
     
     else:
         st.info("👆 Upload both a video and background image to start processing!")
@@ -427,10 +459,9 @@ def update_progress(progress, frame_count, total_frames):
     st.markdown("---")
     st.markdown("""
     <div style="text-align: center; color: #666; padding: 10px;">
-        <p><small>🍹 Powered by MatAnyone and MediaPipe | Optimized for MyAvatar</small></p>
+        <p><small>🍹 Powered by MatAnyone and MediaPipe | Optimized for MyAvatar | API Fixed</small></p>
     </div>
     """, unsafe_allow_html=True)
 
 if __name__ == "__main__":
-    main()
-    
+    main()