Update app.py

app.py

@@ -267,6 +267,73 @@ def process_video_with_green_screen(video_path, background_url, progress_callbac
         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
+
 def main():
     """Streamlit main function"""
     st.set_page_config(
@@ -316,32 +383,70 @@ def main():
     with col2:
         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()),
+        # Background selection method
+        background_method = st.radio(
+            "Choose background method:",
+            ["📋 Preset Backgrounds", "📁 Upload Custom Image"],
             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")
+        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
+                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}")
+                except Exception as e:
+                    st.error(f"Failed to load custom background: {e}")
+                    custom_background = None
+            else:
+                st.info("Please upload a background image")
     
     # Process button
     if uploaded_video 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)
@@ -351,12 +456,19 @@ 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
-            )
+            # 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!")