Update app.py

app.py

@@ -1,98 +1,210 @@
-# HuggingFace Space - Video Background Replacement
-# STREAMLIT VERSION - Memory optimized
+#!/usr/bin/env python3
+"""
+BackgroundFX - Video Background Replacement with Green Screen Workflow
+Hugging Face Space Implementation with SAM2 + MatAnyone
+"""
 
 import streamlit as st
 import cv2
 import numpy as np
-from PIL import Image
 import tempfile
-import traceback
 import os
-import gc
+from PIL import Image
+import requests
+from io import BytesIO
+import logging
 
-# Environment setup
-os.environ["OMP_NUM_THREADS"] = "4"
+# Configure logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
 
-# Page config
-st.set_page_config(
-    page_title="🎬 Video Background Replacement",
-    page_icon="🎬",
-    layout="wide"
-)
+# Try to import SAM2 and MatAnyone
+try:
+    from sam2.build_sam import build_sam2_video_predictor
+    from sam2.sam2_image_predictor import SAM2ImagePredictor
+    SAM2_AVAILABLE = True
+    logger.info("✅ SAM2 loaded successfully")
+except ImportError as e:
+    SAM2_AVAILABLE = False
+    logger.warning(f"⚠️ SAM2 not available: {e}")
 
-# Global model cache
-MODEL_CACHE = {}
+try:
+    import matanyone
+    MATANYONE_AVAILABLE = True
+    logger.info("✅ MatAnyone loaded successfully")
+except ImportError as e:
+    MATANYONE_AVAILABLE = False
+    logger.warning(f"⚠️ MatAnyone not available: {e}")
+
+def load_background_image(background_url):
+    """Load background image from URL"""
+    try:
+        response = requests.get(background_url)
+        response.raise_for_status()
+        image = Image.open(BytesIO(response.content))
+        return np.array(image.convert('RGB'))
+    except Exception as e:
+        logger.error(f"Failed to load background image: {e}")
+        # Return default brick wall background
+        return create_default_background()
 
-def load_models():
-    """Load required models with caching"""
-    global MODEL_CACHE
+def create_default_background():
+    """Create a default brick wall background"""
+    # Create a simple brick pattern
+    height, width = 720, 1280
+    background = np.ones((height, width, 3), dtype=np.uint8) * 150
     
-    if 'models_loaded' in MODEL_CACHE:
-        return True
+    # Add brick pattern
+    brick_height, brick_width = 40, 80
+    for y in range(0, height, brick_height):
+        for x in range(0, width, brick_width):
+            # Alternate brick offset
+            offset = brick_width // 2 if (y // brick_height) % 2 else 0
+            x_pos = (x + offset) % width
+            
+            # Draw brick
+            cv2.rectangle(background, 
+                         (x_pos, y), 
+                         (min(x_pos + brick_width - 2, width), min(y + brick_height - 2, height)), 
+                         (180, 120, 80), -1)
+            cv2.rectangle(background, 
+                         (x_pos, y), 
+                         (min(x_pos + brick_width - 2, width), min(y + brick_height - 2, height)), 
+                         (120, 80, 40), 2)
     
+    return background
+
+def segment_person_sam2(frame):
+    """Segment person using SAM2"""
     try:
-        # Try to load SAM2 if available
-        try:
-            from sam2.sam2_image_predictor import SAM2ImagePredictor
-            predictor = SAM2ImagePredictor.from_pretrained("facebook/sam2-hiera-small")
-            MODEL_CACHE['sam2_predictor'] = predictor
-            st.success("✅ SAM2 model loaded")
-        except ImportError:
-            st.warning("⚠️ SAM2 not available, using fallback method")
-            MODEL_CACHE['sam2_predictor'] = None
+        # Initialize SAM2 predictor
+        predictor = SAM2ImagePredictor.from_pretrained("facebook/sam2-hiera-large")
         
-        MODEL_CACHE['models_loaded'] = True
-        return True
+        # Set image
+        predictor.set_image(frame)
+        
+        # Use center point as prompt (assuming person is in center)
+        h, w = frame.shape[:2]
+        center_point = np.array([[w//2, h//2]])
+        center_label = np.array([1])
+        
+        # Predict mask
+        masks, scores, _ = predictor.predict(
+            point_coords=center_point,
+            point_labels=center_label,
+            multimask_output=False
+        )
+        
+        return masks[0] if len(masks) > 0 else None
         
     except Exception as e:
-        st.error(f"❌ Error loading models: {e}")
-        return False
+        logger.error(f"SAM2 segmentation failed: {e}")
+        return None
 
-def apply_simple_background_replacement(frame, background):
-    """Simple background replacement using color segmentation"""
+def segment_person_fallback(frame):
+    """Fallback person segmentation using color-based method"""
     try:
-        # Convert to HSV
-        hsv = cv2.cvtColor(frame, cv2.COLOR_BGR2HSV)
+        # Convert to HSV for better skin detection
+        hsv = cv2.cvtColor(frame, cv2.COLOR_RGB2HSV)
         
-        # Create mask for person (basic skin tone detection)
+        # Define skin color range
         lower_skin = np.array([0, 20, 70])
         upper_skin = np.array([20, 255, 255])
-        mask1 = cv2.inRange(hsv, lower_skin, upper_skin)
         
-        # Additional mask for different skin tones
-        lower_skin2 = np.array([160, 20, 70])
-        upper_skin2 = np.array([180, 255, 255])
-        mask2 = cv2.inRange(hsv, lower_skin2, upper_skin2)
+        # Create mask for skin tones
+        skin_mask = cv2.inRange(hsv, lower_skin, upper_skin)
         
-        # Combine masks
-        mask = cv2.bitwise_or(mask1, mask2)
+        # Morphological operations to clean up mask
+        kernel = np.ones((5, 5), np.uint8)
+        skin_mask = cv2.morphologyEx(skin_mask, cv2.MORPH_CLOSE, kernel)
+        skin_mask = cv2.morphologyEx(skin_mask, cv2.MORPH_OPEN, kernel)
         
-        # Clean up mask
-        kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (5, 5))
-        mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernel)
-        mask = cv2.GaussianBlur(mask, (5, 5), 0)
+        # Find largest contour (assumed to be person)
+        contours, _ = cv2.findContours(skin_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
         
-        # Apply background replacement
-        mask_norm = mask.astype(np.float32) / 255.0
-        mask_3ch = np.stack([mask_norm] * 3, axis=-1)
+        if contours:
+            # Get largest contour
+            largest_contour = max(contours, key=cv2.contourArea)
+            
+            # Create mask from contour
+            mask = np.zeros(frame.shape[:2], dtype=np.uint8)
+            cv2.fillPoly(mask, [largest_contour], 255)
+            
+            # Expand mask to include more of the person
+            kernel = np.ones((20, 20), np.uint8)
+            mask = cv2.dilate(mask, kernel, iterations=2)
+            
+            return mask.astype(bool)
         
-        result = frame * mask_3ch + background * (1 - mask_3ch)
-        return result.astype(np.uint8)
+        return None
         
     except Exception as e:
-        st.error(f"Error in background replacement: {e}")
+        logger.error(f"Fallback segmentation failed: {e}")
+        return None
+
+def insert_green_screen(frame, person_mask):
+    """Insert green screen background while preserving person"""
+    try:
+        # Create green background
+        green_background = np.zeros_like(frame)
+        green_background[:, :] = [0, 255, 0]  # Pure green (RGB)
+        
+        # Combine person with green background
+        # Where mask is True (person), keep original frame
+        # Where mask is False (background), use green
+        result = np.where(person_mask[..., None], frame, green_background)
+        
+        return result
+        
+    except Exception as e:
+        logger.error(f"Green screen insertion failed: {e}")
         return frame
 
-def process_video(input_video, background_image=None):
-    """Process video with background replacement"""
+def chroma_key_replacement(green_screen_frame, new_background):
+    """Replace green screen with new background using chroma key"""
+    try:
+        # Resize background to match frame
+        h, w = green_screen_frame.shape[:2]
+        background_resized = cv2.resize(new_background, (w, h))
+        
+        # Convert to HSV for better green detection
+        hsv = cv2.cvtColor(green_screen_frame, cv2.COLOR_RGB2HSV)
+        
+        # Define green color range for chroma key
+        lower_green = np.array([40, 50, 50])
+        upper_green = np.array([80, 255, 255])
+        
+        # Create mask for green pixels
+        green_mask = cv2.inRange(hsv, lower_green, upper_green)
+        
+        # Smooth the mask
+        kernel = np.ones((3, 3), np.uint8)
+        green_mask = cv2.morphologyEx(green_mask, cv2.MORPH_CLOSE, kernel)
+        green_mask = cv2.GaussianBlur(green_mask, (5, 5), 0)
+        
+        # Normalize mask to 0-1 range
+        mask_normalized = green_mask.astype(float) / 255
+        
+        # Apply chroma key replacement
+        result = green_screen_frame.copy()
+        for c in range(3):
+            result[:, :, c] = (green_screen_frame[:, :, c] * (1 - mask_normalized) + 
+                              background_resized[:, :, c] * mask_normalized)
+        
+        return result.astype(np.uint8)
+        
+    except Exception as e:
+        logger.error(f"Chroma key replacement failed: {e}")
+        return green_screen_frame
+
+def process_video_with_green_screen(video_path, background_url, progress_callback=None):
+    """Process video with proper green screen workflow"""
     try:
-        if not os.path.exists(input_video):
-            return None, "❌ Input video not found"
+        # Load background image
+        background_image = load_background_image(background_url)
         
-        # Load video
-        cap = cv2.VideoCapture(input_video)
-        if not cap.isOpened():
-            return None, "❌ Could not open video"
+        # Open video
+        cap = cv2.VideoCapture(video_path)
         
         # Get video properties
         fps = int(cap.get(cv2.CAP_PROP_FPS))
@@ -100,144 +212,189 @@ def process_video(input_video, background_image=None):
         height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
         total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
         
-        st.info(f"📊 Video: {width}x{height}, {fps} FPS, {total_frames} frames")
-        
-        # Load background
-        if background_image and os.path.exists(background_image):
-            bg_img = cv2.imread(background_image)
-            bg_img = cv2.resize(bg_img, (width, height))
-        else:
-            # Green screen default
-            bg_img = np.zeros((height, width, 3), dtype=np.uint8)
-            bg_img[:, :] = [0, 255, 0]
-        
-        # Create output video
+        # Create output video writer
         output_path = tempfile.mktemp(suffix='.mp4')
         fourcc = cv2.VideoWriter_fourcc(*'mp4v')
         out = cv2.VideoWriter(output_path, fourcc, fps, (width, height))
         
-        # Progress bar
-        progress_bar = st.progress(0)
-        status_text = st.empty()
-        
         frame_count = 0
-        processed_count = 0
         
-        # Process frames
         while True:
             ret, frame = cap.read()
             if not ret:
                 break
             
-            frame_count += 1
+            # Convert BGR to RGB
+            frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
             
-            # Process frame
-            processed_frame = apply_simple_background_replacement(frame, bg_img)
-            out.write(processed_frame)
-            processed_count += 1
+            # Step 1: Segment person
+            if SAM2_AVAILABLE:
+                person_mask = segment_person_sam2(frame_rgb)
+                method_used = "SAM2"
+            else:
+                person_mask = segment_person_fallback(frame_rgb)
+                method_used = "Fallback"
             
-            # Update progress
-            progress = frame_count / total_frames
-            progress_bar.progress(progress)
-            status_text.text(f"Processing frame {frame_count}/{total_frames}")
+            if person_mask is not None:
+                # Step 2: Insert green screen
+                green_screen_frame = insert_green_screen(frame_rgb, person_mask)
+                
+                # Step 3: Chroma key replacement
+                final_frame = chroma_key_replacement(green_screen_frame, background_image)
+            else:
+                # If segmentation fails, use original frame
+                final_frame = frame_rgb
+                method_used = "No segmentation"
+            
+            # Convert back to BGR for video writer
+            final_frame_bgr = cv2.cvtColor(final_frame, cv2.COLOR_RGB2BGR)
+            out.write(final_frame_bgr)
+            
+            frame_count += 1
             
-            # Memory cleanup every 100 frames
-            if frame_count % 100 == 0:
-                gc.collect()
+            # Update progress
+            if progress_callback:
+                progress = frame_count / total_frames
+                progress_callback(progress, f"Processing frame {frame_count}/{total_frames} ({method_used})")
         
+        # Release resources
         cap.release()
         out.release()
         
-        # Final cleanup
-        gc.collect()
+        return output_path
         
-        if os.path.exists(output_path) and os.path.getsize(output_path) > 0:
-            return output_path, f"✅ Processed {processed_count} frames successfully!"
-        else:
-            return None, "❌ Failed to create output video"
-            
     except Exception as e:
-        return None, f"❌ Error: {str(e)}"
+        logger.error(f"Video processing failed: {e}")
+        return None
 
-# Streamlit UI
 def main():
-    st.title("🎬 AI Video Background Replacement")
-    st.markdown("Upload a video and optionally a background image to replace the background using AI.")
+    """Streamlit main function"""
+    st.set_page_config(
+        page_title="BackgroundFX - Video Background Replacement",
+        page_icon="🎬",
+        layout="wide"
+    )
+    
+    st.title("🎬 BackgroundFX - Video Background Replacement")
+    st.markdown("**Professional video background replacement with green screen workflow**")
+    
+    # Show available methods
+    methods = []
+    if SAM2_AVAILABLE:
+        methods.append("✅ SAM2 (AI Segmentation)")
+    if MATANYONE_AVAILABLE:
+        methods.append("✅ MatAnyone (Advanced Processing)")
+    methods.append("✅ Fallback Method (Color-based)")
     
-    # Load models
-    if not load_models():
-        st.error("Failed to load required models")
-        return
+    st.sidebar.markdown("### Available Methods")
+    for method in methods:
+        st.sidebar.markdown(method)
     
+    # File upload
     col1, col2 = st.columns(2)
     
     with col1:
-        st.subheader("📹 Input Video")
+        st.markdown("### 📹 Upload Video")
         uploaded_video = st.file_uploader(
-            "Choose a video file",
+            "Choose a video file", 
             type=['mp4', 'avi', 'mov', 'mkv'],
             help="Upload the video you want to process"
         )
         
         if uploaded_video:
+            # Display video info
+            st.success(f"✅ Video uploaded: {uploaded_video.name}")
+            
             # Save uploaded video
-            with tempfile.NamedTemporaryFile(delete=False, suffix='.mp4') as tmp_video:
-                tmp_video.write(uploaded_video.read())
-                video_path = tmp_video.name
+            with tempfile.NamedTemporaryFile(delete=False, suffix='.mp4') as tmp_file:
+                tmp_file.write(uploaded_video.read())
+                video_path = tmp_file.name
             
-            st.video(video_path)
+            # Show video preview
+            st.video(uploaded_video)
     
     with col2:
-        st.subheader("🖼️ Background Image")
-        uploaded_bg = st.file_uploader(
-            "Choose background image (optional)",
-            type=['jpg', 'jpeg', 'png'],
-            help="Leave empty for green screen background"
+        st.markdown("### 🖼️ Background Image")
+        
+        # Default background options
+        background_options = {
+            "Brick Wall": "https://images.unsplash.com/photo-1558618666-fcd25c85cd64?w=1280&h=720&fit=crop",
+            "Office": "https://images.unsplash.com/photo-1497366216548-37526070297c?w=1280&h=720&fit=crop",
+            "Nature": "https://images.unsplash.com/photo-1506905925346-21bda4d32df4?w=1280&h=720&fit=crop",
+            "City": "https://images.unsplash.com/photo-1449824913935-59a10b8d2000?w=1280&h=720&fit=crop"
+        }
+        
+        selected_background = st.selectbox(
+            "Choose background",
+            options=list(background_options.keys()),
+            index=0
         )
         
-        bg_path = None
-        if uploaded_bg:
-            # Save uploaded background
-            with tempfile.NamedTemporaryFile(delete=False, suffix='.jpg') as tmp_bg:
-                tmp_bg.write(uploaded_bg.read())
-                bg_path = tmp_bg.name
-            
-            st.image(bg_path, caption="Background Image")
-        else:
-            st.info("💚 Using green screen background")
+        background_url = background_options[selected_background]
+        
+        # Show background preview
+        try:
+            background_image = load_background_image(background_url)
+            st.image(background_image, caption=f"Background: {selected_background}", use_column_width=True)
+        except:
+            st.error("Failed to load background image")
     
     # Process button
-    if st.button("🚀 Process Video", type="primary"):
-        if uploaded_video:
-            with st.spinner("Processing video..."):
-                result_path, message = process_video(video_path, bg_path)
+    if uploaded_video and st.button("🎬 Process Video", type="primary"):
+        
+        with st.spinner("Processing video with green screen workflow..."):
+            # Create progress bar
+            progress_bar = st.progress(0)
+            status_text = st.empty()
+            
+            def update_progress(progress, message):
+                progress_bar.progress(progress)
+                status_text.text(message)
+            
+            # Process video
+            output_path = process_video_with_green_screen(
+                video_path, 
+                background_url, 
+                progress_callback=update_progress
+            )
+            
+            if output_path and os.path.exists(output_path):
+                st.success("✅ Video processing completed!")
+                
+                # Display processed video
+                st.markdown("### 🎉 Processed Video")
                 
-                if result_path:
-                    st.success(message)
-                    st.subheader("✨ Processed Video")
-                    st.video(result_path)
-                    
-                    # Download button
-                    with open(result_path, 'rb') as f:
-                        st.download_button(
-                            label="📥 Download Processed Video",
-                            data=f.read(),
-                            file_name="processed_video.mp4",
-                            mime="video/mp4"
-                        )
-                else:
-                    st.error(message)
-        else:
-            st.warning("Please upload a video first!")
+                with open(output_path, 'rb') as video_file:
+                    video_bytes = video_file.read()
+                    st.video(video_bytes)
+                
+                # Download button
+                st.download_button(
+                    label="📥 Download Processed Video",
+                    data=video_bytes,
+                    file_name=f"backgroundfx_{uploaded_video.name}",
+                    mime="video/mp4"
+                )
+                
+                # Cleanup
+                try:
+                    os.unlink(video_path)
+                    os.unlink(output_path)
+                except:
+                    pass
+            else:
+                st.error("❌ Video processing failed. Please try again.")
     
-    # Status info
+    # Footer
     st.markdown("---")
+    st.markdown("### 🔧 Technical Details")
     st.markdown("""
-    ### 📊 Current Status:
-    - ✅ **Video processing**: Active
-    - 🎨 **Background replacement**: Working  
-    - 💾 **Memory optimization**: Enabled
-    - 🔄 **Fallback processing**: Available
+    **Green Screen Workflow:**
+    1. **Person Segmentation** - AI identifies the person in each frame
+    2. **Green Screen Insert** - Replaces background with pure green
+    3. **Chroma Key Replacement** - Replaces green with new background
+    
+    This ensures clean edges and professional results.
     """)
 
 if __name__ == "__main__":