Update streamlit_app.py

streamlit_app.py

@@ -1,4 +1,13 @@
-# streamlit_ui.py - Simplified single-button workflow
+# ==================================================================================
+# STREAMLIT VIDEO BACKGROUND REPLACER - MAIN APPLICATION
+# ==================================================================================
+# Single-button workflow: Upload video + background → Process → Download result
+# Uses SAM2 + MatAnyone pipeline with temporal smoothing
+# ==================================================================================
+
+# ==================================================================================
+# CHAPTER 1: IMPORTS AND SETUP
+# ==================================================================================
 
 import streamlit as st
 import os
@@ -28,6 +37,10 @@ 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
 
+# ==================================================================================
+# CHAPTER 2: STREAMLIT CONFIGURATION
+# ==================================================================================
+
 # Streamlit Page Config
 st.set_page_config(
     page_title="Advanced Video Background Replacer",
@@ -36,7 +49,30 @@ def custom_excepthook(type, value, tb):
     initial_sidebar_state="expanded"
 )
 
-# GPU Diagnostic
+# Custom CSS
+st.markdown("""
+<style>
+    .main .block-container {
+        padding-top: 2rem;
+        padding-bottom: 2rem;
+    }
+    .stButton>button {
+        width: 100%;
+        background-color: #4CAF50;
+        color: white;
+        font-weight: bold;
+        transition: all 0.3s;
+    }
+    .stButton>button:hover {
+        background-color: #45a049;
+    }
+</style>
+""", unsafe_allow_html=True)
+
+# ==================================================================================
+# CHAPTER 3: GPU DIAGNOSTICS
+# ==================================================================================
+
 def check_gpu():
     """Check GPU availability and log details"""
     logger.info("=" * 60)
@@ -67,27 +103,10 @@ def check_gpu():
     logger.info("=" * 60)
     return cuda_available
 
-# Custom CSS
-st.markdown("""
-<style>
-    .main .block-container {
-        padding-top: 2rem;
-        padding-bottom: 2rem;
-    }
-    .stButton>button {
-        width: 100%;
-        background-color: #4CAF50;
-        color: white;
-        font-weight: bold;
-        transition: all 0.3s;
-    }
-    .stButton>button:hover {
-        background-color: #45a049;
-    }
-</style>
-""", unsafe_allow_html=True)
+# ==================================================================================
+# CHAPTER 4: SESSION STATE INITIALIZATION
+# ==================================================================================
 
-# Session State Initialization
 def initialize_session_state():
     """Initialize all session state variables"""
     defaults = {
@@ -113,13 +132,20 @@ def initialize_session_state():
     if st.session_state.gpu_available is None:
         st.session_state.gpu_available = check_gpu()
 
-# Main Application
+# ==================================================================================
+# CHAPTER 5: MAIN APPLICATION
+# ==================================================================================
+
 def main():
     st.title("Advanced Video Background Replacer")
     st.markdown("---")
     
     initialize_session_state()
     
+    # ==================================================================================
+    # SIDEBAR: System Status
+    # ==================================================================================
+    
     with st.sidebar:
         st.subheader("System Status")
         if st.session_state.gpu_available:
@@ -127,7 +153,10 @@ def main():
         else:
             st.error("GPU: Not Available")
     
-    # Display last error if exists
+    # ==================================================================================
+    # ERROR DISPLAY
+    # ==================================================================================
+    
     if st.session_state.last_error:
         with st.expander("⚠️ Last Error", expanded=False):
             st.error(st.session_state.last_error)
@@ -137,6 +166,10 @@ def main():
     
     col1, col2 = st.columns([1, 1], gap="large")
     
+    # ==================================================================================
+    # COLUMN 1: VIDEO UPLOAD
+    # ==================================================================================
+    
     with col1:
         st.header("1. Upload Video")
         
@@ -146,10 +179,10 @@ def main():
             key="video_uploader"
         )
         
-        # Check if video changed using id()
+        # FIX: Only update session state when uploaded is not None
         current_video_id = id(uploaded)
-        if current_video_id != st.session_state.last_video_id:
-            logger.info(f"New video: {uploaded.name if uploaded else 'None'}")
+        if uploaded is not None and current_video_id != st.session_state.last_video_id:
+            logger.info(f"New video: {uploaded.name}")
             st.session_state.uploaded_video = uploaded
             st.session_state.last_video_id = current_video_id
             st.session_state.video_bytes_cache = None
@@ -179,6 +212,10 @@ def main():
         else:
             st.session_state.video_preview_placeholder.empty()
     
+    # ==================================================================================
+    # COLUMN 2: BACKGROUND SETTINGS & PROCESSING
+    # ==================================================================================
+    
     with col2:
         st.header("2. Background Settings")
         
@@ -189,6 +226,10 @@ def main():
             key="bg_type_radio"
         )
         
+        # ==================================================================================
+        # BACKGROUND IMAGE UPLOAD
+        # ==================================================================================
+        
         if bg_type == "Image":
             bg_image = st.file_uploader(
                 "Upload Background Image",
@@ -196,15 +237,12 @@ def main():
                 key="bg_image_uploader"
             )
             
-            # Check if image changed using id()
+            # FIX: Only update session state when bg_image is not None
             current_bg_id = id(bg_image)
-            if current_bg_id != st.session_state.last_bg_image_id:
-                logger.info(f"New background: {bg_image.name if bg_image else 'None'}")
+            if bg_image is not None and current_bg_id != st.session_state.last_bg_image_id:
+                logger.info(f"New background: {bg_image.name}")
                 st.session_state.last_bg_image_id = current_bg_id
-                if bg_image is not None:
-                    st.session_state.bg_image_cache = Image.open(bg_image)
-                else:
-                    st.session_state.bg_image_cache = None
+                st.session_state.bg_image_cache = Image.open(bg_image)
             
             # Background preview with anti-shake placeholder
             if st.session_state.bg_preview_placeholder is None:
@@ -216,6 +254,10 @@ def main():
             else:
                 st.session_state.bg_preview_placeholder.empty()
         
+        # ==================================================================================
+        # BACKGROUND COLOR PICKER
+        # ==================================================================================
+        
         elif bg_type == "Color":
             selected_color = st.color_picker(
                 "Choose Background Color",
@@ -242,6 +284,10 @@ def main():
             else:
                 st.session_state.bg_preview_placeholder.empty()
         
+        # ==================================================================================
+        # VIDEO PROCESSING SECTION
+        # ==================================================================================
+        
         st.header("3. Process Video")
         
         can_process = (
@@ -308,7 +354,10 @@ def main():
             finally:
                 st.session_state.processing = False
         
-        # Show processed video if available
+        # ==================================================================================
+        # RESULTS DISPLAY
+        # ==================================================================================
+        
         if st.session_state.processed_video_bytes is not None:
             st.markdown("---")
             st.markdown("### Processed Video")
@@ -326,5 +375,9 @@ def main():
                 logger.error(f"Display error: {e}")
                 st.error(f"Display error: {e}")
 
+# ==================================================================================
+# CHAPTER 6: APPLICATION ENTRY POINT
+# ==================================================================================
+
 if __name__ == "__main__":
     main()