Update app.py

app.py

@@ -1,7 +1,8 @@
 #!/usr/bin/env python3
 """
 BackgroundFX - Video Background Replacement with Green Screen Workflow
-Hugging Face Space Implementation with SAM2 + MatAnyone
+Fixed for Hugging Face Space - Handles video preview issues
+FIXED: Video display issue by properly handling file stream
 """
 
 import streamlit as st
@@ -13,6 +14,7 @@
 import requests
 from io import BytesIO
 import logging
+import base64
 
 # Configure logging
 logging.basicConfig(level=logging.INFO)
@@ -36,6 +38,41 @@
     MATANYONE_AVAILABLE = False
     logger.warning(f"⚠️ MatAnyone not available: {e}")
 
+def get_video_info(video_path):
+    """Get video information and first frame"""
+    try:
+        cap = cv2.VideoCapture(video_path)
+        if not cap.isOpened():
+            return None, None
+        
+        # Get video properties
+        fps = int(cap.get(cv2.CAP_PROP_FPS))
+        width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
+        height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
+        total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
+        duration = total_frames / fps if fps > 0 else 0
+        
+        # Get first frame for thumbnail
+        ret, first_frame = cap.read()
+        cap.release()
+        
+        if ret:
+            # Convert BGR to RGB
+            first_frame_rgb = cv2.cvtColor(first_frame, cv2.COLOR_BGR2RGB)
+            return {
+                'fps': fps,
+                'width': width,
+                'height': height,
+                'total_frames': total_frames,
+                'duration': duration
+            }, first_frame_rgb
+        
+        return None, None
+        
+    except Exception as e:
+        logger.error(f"Error getting video info: {e}")
+        return None, None
+
 def load_background_image(background_url):
     """Load background image from URL"""
     try:
@@ -45,12 +82,10 @@ def load_background_image(background_url):
         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 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
     
@@ -58,11 +93,9 @@ def create_default_background():
     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)), 
@@ -77,18 +110,13 @@ def create_default_background():
 def segment_person_sam2(frame):
     """Segment person using SAM2"""
     try:
-        # Initialize SAM2 predictor
         predictor = SAM2ImagePredictor.from_pretrained("facebook/sam2-hiera-large")
-        
-        # 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,
@@ -104,36 +132,25 @@ def segment_person_sam2(frame):
 def segment_person_fallback(frame):
     """Fallback person segmentation using color-based method"""
     try:
-        # Convert to HSV for better skin detection
         hsv = cv2.cvtColor(frame, cv2.COLOR_RGB2HSV)
         
-        # Define skin color range
         lower_skin = np.array([0, 20, 70])
         upper_skin = np.array([20, 255, 255])
         
-        # Create mask for skin tones
         skin_mask = cv2.inRange(hsv, lower_skin, upper_skin)
         
-        # 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)
         
-        # Find largest contour (assumed to be person)
         contours, _ = cv2.findContours(skin_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
         
         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)
         
         return None
@@ -145,65 +162,43 @@ def segment_person_fallback(frame):
 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)
+        green_background[:, :] = [0, 255, 0]  # Pure green
         
-        # 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 chroma_key_replacement(green_screen_frame, new_background):
-    """Replace green screen with new background using chroma key"""
+def replace_background(frame, person_mask, background):
+    """Replace background with new image"""
     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)
+        # Resize background to match frame dimensions
+        h, w = frame.shape[:2]
+        background_resized = cv2.resize(background, (w, h))
         
-        # Normalize mask to 0-1 range
-        mask_normalized = green_mask.astype(float) / 255
+        # Create smooth edges
+        mask_float = person_mask.astype(np.float32)
+        kernel = np.ones((5, 5), np.float32) / 25
+        mask_float = cv2.filter2D(mask_float, -1, kernel)
         
-        # Apply chroma key replacement
-        result = green_screen_frame.copy()
+        # Composite the images
+        result = np.zeros_like(frame)
         for c in range(3):
-            result[:, :, c] = (green_screen_frame[:, :, c] * (1 - mask_normalized) + 
-                              background_resized[:, :, c] * mask_normalized)
+            result[:, :, c] = (mask_float * frame[:, :, c] + 
+                              (1 - mask_float) * background_resized[:, :, c])
         
         return result.astype(np.uint8)
         
     except Exception as e:
-        logger.error(f"Chroma key replacement failed: {e}")
-        return green_screen_frame
+        logger.error(f"Background replacement failed: {e}")
+        return frame
 
-def process_video_with_green_screen(video_path, background_url, progress_callback=None):
-    """Process video with proper green screen workflow"""
+def process_video(video_path, output_path, background_image, use_green_screen=True, progress_callback=None):
+    """Process video with background replacement"""
     try:
-        # Load background image
-        background_image = load_background_image(background_url)
-        
-        # Open video
         cap = cv2.VideoCapture(video_path)
         
         # Get video properties
@@ -212,8 +207,7 @@ def process_video_with_green_screen(video_path, background_url, progress_callbac
         height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
         total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
         
-        # Create output video writer
-        output_path = tempfile.mktemp(suffix='.mp4')
+        # Setup video writer
         fourcc = cv2.VideoWriter_fourcc(*'mp4v')
         out = cv2.VideoWriter(output_path, fourcc, fps, (width, height))
         
@@ -224,296 +218,230 @@ def process_video_with_green_screen(video_path, background_url, progress_callbac
             if not ret:
                 break
             
-            # Convert BGR to RGB
+            # Convert BGR to RGB for processing
             frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
             
-            # Step 1: Segment person
+            # 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"
             
             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"
+                if use_green_screen:
+                    # First create green screen
+                    frame_rgb = insert_green_screen(frame_rgb, person_mask)
+                    # Then replace green with background
+                    # Create green mask
+                    lower_green = np.array([0, 200, 0])
+                    upper_green = np.array([100, 255, 100])
+                    green_mask = cv2.inRange(frame_rgb, lower_green, upper_green)
+                    green_mask = green_mask.astype(bool)
+                    # Replace green areas with background
+                    frame_rgb = replace_background(frame_rgb, ~green_mask, background_image)
+                else:
+                    # Direct background replacement
+                    frame_rgb = replace_background(frame_rgb, person_mask, background_image)
             
-            # Convert back to BGR for video writer
-            final_frame_bgr = cv2.cvtColor(final_frame, cv2.COLOR_RGB2BGR)
-            out.write(final_frame_bgr)
+            # Convert back to BGR for writing
+            frame_bgr = cv2.cvtColor(frame_rgb, cv2.COLOR_RGB2BGR)
+            out.write(frame_bgr)
             
             frame_count += 1
-            
-            # Update progress
             if progress_callback:
-                progress = frame_count / total_frames
-                progress_callback(progress, f"Processing frame {frame_count}/{total_frames} ({method_used})")
+                progress_callback(frame_count / total_frames)
         
-        # Release resources
         cap.release()
         out.release()
         
-        return output_path
+        return True
         
     except Exception as e:
         logger.error(f"Video processing failed: {e}")
-        return None
-
-def process_video_with_green_screen_custom(video_path, background_image, progress_callback=None):
-    """Process video with custom background image"""
-    try:
-        # Open video
-        cap = cv2.VideoCapture(video_path)
-        
-        # Get video properties
-        fps = int(cap.get(cv2.CAP_PROP_FPS))
-        width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
-        height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
-        total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
-        
-        # Create output video writer
-        output_path = tempfile.mktemp(suffix='.mp4')
-        fourcc = cv2.VideoWriter_fourcc(*'mp4v')
-        out = cv2.VideoWriter(output_path, fourcc, fps, (width, height))
-        
-        frame_count = 0
-        
-        while True:
-            ret, frame = cap.read()
-            if not ret:
-                break
-            
-            # Convert BGR to RGB
-            frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
-            
-            # 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"
-            
-            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 with custom background
-                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
-            
-            # 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()
-        
-        return output_path
-        
-    except Exception as e:
-        logger.error(f"Video processing with custom background failed: {e}")
-        return None
+        return False
 
+# Streamlit UI
 def main():
-    """Streamlit main function"""
     st.set_page_config(
-        page_title="BackgroundFX - Video Background Replacement",
+        page_title="BackgroundFX",
         page_icon="🎬",
         layout="wide"
     )
     
     st.title("🎬 BackgroundFX - Video Background Replacement")
-    st.markdown("**Professional video background replacement with green screen workflow**")
+    st.markdown("Replace video backgrounds with AI-powered segmentation")
     
-    # 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)")
+    # Check dependencies
+    if not SAM2_AVAILABLE:
+        st.warning("⚠️ SAM2 not available - using fallback segmentation method")
+    if not MATANYONE_AVAILABLE:
+        st.info("ℹ️ MatAnyone not available - using standard matting")
     
-    st.sidebar.markdown("### Available Methods")
-    for method in methods:
-        st.sidebar.markdown(method)
+    # Sidebar for settings
+    with st.sidebar:
+        st.header("⚙️ Settings")
+        
+        use_green_screen = st.checkbox(
+            "Use Green Screen Workflow",
+            value=True,
+            help="First create green screen, then replace with background"
+        )
+        
+        st.subheader("📸 Background Options")
+        bg_option = st.radio(
+            "Choose background source:",
+            ["Default Brick Wall", "Custom URL", "Upload Image"]
+        )
+        
+        background_image = None
+        
+        if bg_option == "Custom URL":
+            bg_url = st.text_input(
+                "Background Image URL",
+                value="https://images.unsplash.com/photo-1557683316-973673baf926",
+                help="Enter a direct image URL"
+            )
+            if bg_url:
+                with st.spinner("Loading background..."):
+                    background_image = load_background_image(bg_url)
+                    if background_image is not None:
+                        st.success("✅ Background loaded")
+                        st.image(background_image, caption="Background Preview", use_column_width=True)
+        
+        elif bg_option == "Upload Image":
+            uploaded_bg = st.file_uploader(
+                "Upload Background Image",
+                type=['jpg', 'jpeg', 'png'],
+                help="Upload your own background image"
+            )
+            if uploaded_bg is not None:
+                background_image = np.array(Image.open(uploaded_bg).convert('RGB'))
+                st.success("✅ Background uploaded")
+                st.image(background_image, caption="Background Preview", use_column_width=True)
+        
+        else:  # Default Brick Wall
+            background_image = create_default_background()
+            st.info("Using default brick wall background")
+            st.image(background_image, caption="Default Background", use_column_width=True)
     
-    # File upload
+    # Main content area
     col1, col2 = st.columns(2)
     
-    # Initialize video_path variable
-    video_path = None
-    
     with col1:
-        st.markdown("### 📹 Upload Video")
+        st.header("📹 Input Video")
         uploaded_video = st.file_uploader(
-            "Choose a video file", 
+            "Upload your video",
             type=['mp4', 'avi', 'mov', 'mkv'],
-            help="Upload the video you want to process"
+            help="Upload a video to process"
         )
         
-        if uploaded_video:
-            # Display video info
-            st.success(f"✅ Video uploaded: {uploaded_video.name}")
-            
-            # Read video data once and store it
+        if uploaded_video is not None:
+            # FIXED: Read bytes once and reuse
             video_bytes = uploaded_video.read()
             
-            # Save uploaded video
+            # Save to temporary file for processing
             with tempfile.NamedTemporaryFile(delete=False, suffix='.mp4') as tmp_file:
                 tmp_file.write(video_bytes)
-                video_path = tmp_file.name
+                temp_video_path = tmp_file.name
             
-            # Show video preview using the bytes
+            # Display video using the original bytes (not consuming the stream)
             st.video(video_bytes)
+            
+            # Get video info using the temp file
+            video_info, first_frame = get_video_info(temp_video_path)
+            
+            if video_info and first_frame is not None:
+                st.success(f"✅ Video loaded: {video_info['width']}x{video_info['height']}, "
+                          f"{video_info['fps']} fps, {video_info['duration']:.1f}s")
+                
+                # Show first frame as thumbnail
+                st.image(first_frame, caption="First Frame Preview", use_column_width=True)
+                
+                # Store paths in session state
+                if 'temp_video_path' not in st.session_state:
+                    st.session_state.temp_video_path = temp_video_path
+            else:
+                st.error("Failed to read video information")
     
     with col2:
-        st.markdown("### 🖼️ Background Image")
-        
-        # Background selection method
-        background_method = st.radio(
-            "Choose background method:",
-            ["📋 Preset Backgrounds", "📁 Upload Custom Image"],
-            index=0
-        )
+        st.header("🎯 Output")
         
-        background_url = None
-        custom_background = None
-        
-        if background_method == "📋 Preset Backgrounds":
-            # 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
-            )
-            
-            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")
-        
-        else:  # Upload Custom Image
-            uploaded_background = st.file_uploader(
-                "Upload your background image",
-                type=['jpg', 'jpeg', 'png', 'bmp'],
-                help="Upload a custom background image (JPG, PNG, BMP)"
-            )
-            
-            if uploaded_background:
-                # Load and display custom background
+        if uploaded_video is not None and background_image is not None:
+            if st.button("🚀 Process Video", type="primary"):
                 try:
-                    custom_background = np.array(Image.open(uploaded_background).convert('RGB'))
-                    st.image(custom_background, caption="Custom Background", use_column_width=True)
-                    st.success(f"✅ Custom background uploaded: {uploaded_background.name}")
+                    # Create output path
+                    output_path = tempfile.mktemp(suffix='.mp4')
+                    
+                    # Process video with progress bar
+                    progress_bar = st.progress(0)
+                    status_text = st.empty()
+                    
+                    def update_progress(progress):
+                        progress_bar.progress(progress)
+                        status_text.text(f"Processing: {int(progress * 100)}%")
+                    
+                    status_text.text("Starting video processing...")
+                    
+                    # Use the temp file path we saved
+                    success = process_video(
+                        st.session_state.temp_video_path,
+                        output_path,
+                        background_image,
+                        use_green_screen,
+                        update_progress
+                    )
+                    
+                    if success and os.path.exists(output_path):
+                        status_text.text("✅ Processing complete!")
+                        
+                        # Read the processed video
+                        with open(output_path, 'rb') as f:
+                            processed_video = f.read()
+                        
+                        # Display processed video
+                        st.video(processed_video)
+                        
+                        # Download button
+                        st.download_button(
+                            label="📥 Download Processed Video",
+                            data=processed_video,
+                            file_name="backgroundfx_output.mp4",
+                            mime="video/mp4"
+                        )
+                        
+                        # Cleanup
+                        os.unlink(output_path)
+                    else:
+                        st.error("❌ Video processing failed")
+                    
                 except Exception as e:
-                    st.error(f"Failed to load custom background: {e}")
-                    custom_background = None
-            else:
-                st.info("Please upload a background image")
+                    st.error(f"Error during processing: {str(e)}")
+                    logger.error(f"Processing error: {e}")
+        
+        elif uploaded_video is None:
+            st.info("👈 Please upload a video to begin")
+        elif background_image is None:
+            st.info("👈 Please select or upload a background image")
     
-    # Process button
-    if uploaded_video and video_path and st.button("🎬 Process Video", type="primary"):
-        
-        # Check if background is selected
-        if background_method == "📋 Preset Backgrounds" and not background_url:
-            st.error("Please select a preset background")
-            return
-        elif background_method == "📁 Upload Custom Image" and custom_background is None:
-            st.error("Please upload a custom background image")
-            return
-        
-        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 with appropriate background
-            if background_method == "📋 Preset Backgrounds":
-                output_path = process_video_with_green_screen(
-                    video_path, 
-                    background_url, 
-                    progress_callback=update_progress
-                )
-            else:  # Custom background
-                output_path = process_video_with_green_screen_custom(
-                    video_path, 
-                    custom_background, 
-                    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")
-                
-                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.")
+    # Cleanup temporary files on session end
+    if 'temp_video_path' in st.session_state and os.path.exists(st.session_state.temp_video_path):
+        try:
+            os.unlink(st.session_state.temp_video_path)
+            del st.session_state.temp_video_path
+        except:
+            pass
     
     # Footer
     st.markdown("---")
-    st.markdown("### 🔧 Technical Details")
-    st.markdown("""
-    **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.
-    """)
+    st.markdown(
+        """
+        <div style='text-align: center'>
+            <p>BackgroundFX v1.0 | AI-Powered Video Background Replacement</p>
+            <p>Using SAM2 for person segmentation and green screen technology</p>
+        </div>
+        """,
+        unsafe_allow_html=True
+    )
 
 if __name__ == "__main__":
     main()