Update streamlit_app.py

streamlit_app.py

@@ -1,4 +1,6 @@
 # streamlit_ui.py
+
+# --- Imports ---
 import streamlit as st
 import os
 import sys
@@ -10,15 +12,24 @@
 from PIL import Image
 import logging
 import io
+import torch
+import traceback
+from concurrent.futures import ThreadPoolExecutor
 
-# Add project root to path
+# --- Project Setup ---
 sys.path.append(str(Path(__file__).parent.absolute()))
 
-# Configure logging
-logging.basicConfig(level=logging.INFO)
+# --- Logging Configuration ---
+logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
 logger = logging.getLogger(__name__)
 
-# Set page config
+# --- Global Exception Hook ---
+def custom_excepthook(type, value, tb):
+    logger.error(f"Unhandled: {type.__name__}: {value}\n{''.join(traceback.format_tb(tb))}", exc_info=True)
+sys.excepthook = custom_excepthook
+threading.excepthook = custom_excepthook  # Python 3.10+
+
+# --- Streamlit Page Config ---
 st.set_page_config(
     page_title="🎬 Advanced Video Background Replacer",
     page_icon="🎥",
@@ -26,7 +37,7 @@
     initial_sidebar_state="expanded"
 )
 
-# Custom CSS for better UI
+# --- Custom CSS ---
 st.markdown("""
 <style>
     .main .block-container {
@@ -73,6 +84,7 @@
 </style>
 """, unsafe_allow_html=True)
 
+# --- Session State Initialization ---
 def initialize_session_state():
     """Initialize all session state variables"""
     if 'uploaded_video' not in st.session_state:
@@ -90,16 +102,18 @@ def initialize_session_state():
     if 'progress_text' not in st.session_state:
         st.session_state.progress_text = "Ready"
 
+# --- Video Upload Handling ---
 def handle_video_upload():
     """Handle video file upload"""
     uploaded = st.file_uploader(
-        "📹 Upload Video", 
+        "📹 Upload Video",
         type=["mp4", "mov", "avi"],
         key="video_uploader"
     )
     if uploaded is not None:
         st.session_state.uploaded_video = uploaded
 
+# --- Video Preview ---
 def show_video_preview():
     """Show video preview in the UI"""
     st.markdown("### Video Preview")
@@ -108,6 +122,7 @@ def show_video_preview():
         st.video(video_bytes)
         st.session_state.uploaded_video.seek(0)
 
+# --- Background Selection ---
 def handle_background_selection():
     """Handle background selection UI"""
     st.markdown("### Background Options")
@@ -127,14 +142,14 @@ def handle_background_selection():
         if bg_image is not None:
             st.session_state.bg_image = Image.open(bg_image)
             st.image(
-                st.session_state.bg_image, 
-                caption="Selected Background", 
+                st.session_state.bg_image,
+                caption="Selected Background",
                 use_container_width=True
             )
     
     elif bg_type == "Color":
         st.session_state.bg_color = st.color_picker(
-            "🎨 Choose Background Color", 
+            "🎨 Choose Background Color",
             st.session_state.bg_color
         )
         color_rgb = tuple(int(st.session_state.bg_color.lstrip('#')[i:i+2], 16) for i in (0, 2, 4))
@@ -144,6 +159,8 @@ def handle_background_selection():
     
     return bg_type.lower()
 
+# --- Video Processing ---
+@st.fragment(run_if=lambda: st.session_state.processing)
 def process_video(input_file, background, bg_type="image"):
     """
     Process video with the selected background using SAM2 and MatAnyone pipeline.
@@ -156,8 +173,12 @@ def process_video(input_file, background, bg_type="image"):
             
             # Save the uploaded video to a temporary file
             input_path = str(temp_dir / "input.mp4")
+            logger.info(f"Writing video to {input_path}")
             with open(input_path, "wb") as f:
-                f.write(input_file.getvalue())
+                written = f.write(input_file.getvalue())
+                logger.info(f"Wrote {written/1e6:.2f}MB")
+            if not os.path.exists(input_path):
+                raise FileNotFoundError(f"Input video not saved: {input_path}")
             
             # Prepare background
             bg_path = None
@@ -173,43 +194,53 @@ def process_video(input_file, background, bg_type="image"):
                 bg_path = str(temp_dir / "background.jpg")
                 cv2.imwrite(bg_path, np.ones((100, 100, 3), dtype=np.uint8) * color_rgb[::-1])
             
-            # Set up progress callback
-            progress_bar = st.progress(0)
-            status_text = st.empty()
+            # Set up progress placeholders
+            progress_placeholder = st.empty()
+            status_placeholder = st.empty()
             
             def progress_callback(progress, message):
                 progress = max(0, min(1, float(progress)))
-                progress_bar.progress(progress)
-                status_text.text(f"Status: {message}")
-                st.session_state.progress = int(progress * 100)
-                st.session_state.progress_text = message
+                progress_placeholder.progress(progress)
+                status_placeholder.text(f"Status: {message}")
+            
+            # Log GPU state
+            logger.info(f"CUDA Available: {torch.cuda.is_available()}, Device: {torch.cuda.get_device_name(0)}")
+            logger.info(f"GPU Mem Before: {torch.cuda.memory_allocated()/1e9:.2f}GB")
             
             # Process the video
             output_path = str(temp_dir / "output.mp4")
-            
-            # Mock click points (center of the frame)
             click_points = [[0.5, 0.5]]
-            
-            # Import the pipeline processor
             from pipeline.integrated_pipeline import TwoStageProcessor
             
-            # Initialize the processor
-            processor = TwoStageProcessor(temp_dir=str(temp_dir))
+            # Cache processor
+            @st.cache_resource
+            def load_processor(temp_dir):
+                return TwoStageProcessor(temp_dir=temp_dir)
+            processor = load_processor(str(temp_dir))
             
-            # Process the video
-            success = processor.process_video(
-                input_video=input_path,
-                background_video=bg_path if bg_type == "image" else "",
-                click_points=click_points,
-                output_path=output_path,
-                use_matanyone=True,
-                progress_callback=progress_callback
-            )
+            # Run in thread to catch exceptions
+            with ThreadPoolExecutor(max_workers=1) as executor:
+                future = executor.submit(
+                    processor.process_video,
+                    input_video=input_path,
+                    background_video=bg_path if bg_type == "image" else "",
+                    click_points=click_points,
+                    output_path=output_path,
+                    use_matanyone=True,
+                    progress_callback=progress_callback
+                )
+                try:
+                    success = future.result(timeout=300)  # 5min timeout
+                except Exception as e:
+                    logger.error(f"Pipeline thread fail: {traceback.format_exc()}", exc_info=True)
+                    raise
+            
+            logger.info(f"GPU Mem After: {torch.cuda.memory_allocated()/1e9:.2f}GB")
+            torch.cuda.empty_cache()  # Free memory
             
             if not success:
                 raise RuntimeError("Video processing failed")
             
-            # Return the path to the processed video
             return output_path
             
     except Exception as e:
@@ -217,6 +248,7 @@ def progress_callback(progress, message):
         st.error(f"An error occurred during processing: {str(e)}")
         return None
 
+# --- Main Application ---
 def main():
     st.title("🎬 Advanced Video Background Replacer")
     st.markdown("---")
@@ -235,7 +267,6 @@ def main():
     with col2:
         st.header("2. Background Settings")
         bg_type = handle_background_selection()
-
         st.header("3. Process & Download")
         # ------- FORM WRAP START --------
         with st.form("process_form"):
@@ -243,7 +274,6 @@ def main():
                 "🚀 Process Video",
                 disabled=not st.session_state.uploaded_video or st.session_state.processing
             )
-
             process_status = st.empty()
             if submitted and not st.session_state.processing:
                 st.session_state.processing = True
@@ -255,27 +285,23 @@ def main():
                             background = st.session_state.bg_image
                         elif bg_type == "color" and 'bg_color' in st.session_state:
                             background = st.session_state.bg_color
-
                         # Process the video
                         output_path = process_video(
                             st.session_state.uploaded_video,
                             background,
                             bg_type=bg_type
                         )
-
                         if output_path and os.path.exists(output_path):
                             st.session_state.processed_video_path = output_path
                             process_status.success("✅ Video processing complete!")
                         else:
                             process_status.error("❌ Failed to process video. Please check the logs for details.")
-
                     except Exception as e:
                         process_status.error(f"❌ An error occurred: {str(e)}")
                         logger.exception("Video processing failed")
                     finally:
                         st.session_state.processing = False
         # ------- FORM WRAP END --------
-
         # Show processed video if available (outside form, stable)
         if 'processed_video_path' in st.session_state and st.session_state.processed_video_path:
             st.markdown("### Processed Video")
@@ -294,5 +320,6 @@ def main():
                 st.error(f"Error displaying video: {str(e)}")
                 logger.error(f"Error displaying video: {str(e)}", exc_info=True)
 
+# --- Entry Point ---
 if __name__ == "__main__":
-    main()
+    main()