Update streamlit_app.py

streamlit_app.py

@@ -85,20 +85,23 @@ def custom_excepthook(type, value, tb):
 # --- Session State Initialization ---
 def initialize_session_state():
     """Initialize all session state variables"""
-    if 'uploaded_video' not in st.session_state:
-        st.session_state.uploaded_video = None
-    if 'bg_image' not in st.session_state:
-        st.session_state.bg_image = None
-    if 'bg_color' not in st.session_state:
-        st.session_state.bg_color = "#00FF00"
-    if 'processed_video_path' not in st.session_state:
-        st.session_state.processed_video_path = None
-    if 'processing' not in st.session_state:
-        st.session_state.processing = False
-    if 'progress' not in st.session_state:
-        st.session_state.progress = 0
-    if 'progress_text' not in st.session_state:
-        st.session_state.progress_text = "Ready"
+    defaults = {
+        'uploaded_video': None,
+        'video_bytes_cache': None,
+        'bg_image': None,
+        'bg_image_cache': None,
+        'bg_image_name': None,
+        'bg_color': "#00FF00",
+        'cached_color': None,
+        'color_display_cache': None,
+        'processed_video_path': None,
+        'processing': False,
+        'progress': 0,
+        'progress_text': "Ready"
+    }
+    for key, value in defaults.items():
+        if key not in st.session_state:
+            st.session_state[key] = value
 
 # --- Video Upload Handling ---
 def handle_video_upload():
@@ -108,17 +111,21 @@ def handle_video_upload():
         type=["mp4", "mov", "avi"],
         key="video_uploader"
     )
-    if uploaded is not None:
+    # Only update if actually changed
+    if uploaded is not None and uploaded != st.session_state.uploaded_video:
         st.session_state.uploaded_video = uploaded
+        st.session_state.video_bytes_cache = None  # Clear cache when new video uploaded
 
 # --- Video Preview ---
 def show_video_preview():
     """Show video preview in the UI"""
     st.markdown("### Video Preview")
     if st.session_state.uploaded_video is not None:
-        video_bytes = st.session_state.uploaded_video.getvalue()
-        st.video(video_bytes)
-        st.session_state.uploaded_video.seek(0)
+        # Cache video bytes in session state to prevent reloading
+        if st.session_state.video_bytes_cache is None:
+            st.session_state.video_bytes_cache = st.session_state.uploaded_video.getvalue()
+            st.session_state.uploaded_video.seek(0)
+        st.video(st.session_state.video_bytes_cache)
 
 # --- Background Selection ---
 def handle_background_selection():
@@ -128,7 +135,8 @@ def handle_background_selection():
         "Select Background Type:",
         ["Image", "Color", "Blur"],
         horizontal=True,
-        index=0
+        index=0,
+        key="bg_type_radio"
     )
     
     if bg_type == "Image":
@@ -137,10 +145,17 @@ def handle_background_selection():
             type=["jpg", "png", "jpeg"],
             key="bg_image_uploader"
         )
+        # Only process if new image uploaded
         if bg_image is not None:
-            st.session_state.bg_image = Image.open(bg_image)
+            if (st.session_state.bg_image_cache is None or 
+                st.session_state.get('bg_image_name') != bg_image.name):
+                st.session_state.bg_image = Image.open(bg_image)
+                st.session_state.bg_image_cache = st.session_state.bg_image
+                st.session_state.bg_image_name = bg_image.name
+            
+            # Display cached image
             st.image(
-                st.session_state.bg_image,
+                st.session_state.bg_image_cache,
                 caption="Selected Background",
                 use_container_width=True
             )
@@ -148,12 +163,19 @@ def handle_background_selection():
     elif bg_type == "Color":
         st.session_state.bg_color = st.color_picker(
             "🎨 Choose Background Color",
-            st.session_state.bg_color
+            st.session_state.bg_color,
+            key="color_picker"
         )
-        color_rgb = tuple(int(st.session_state.bg_color.lstrip('#')[i:i+2], 16) for i in (0, 2, 4))
-        color_display = np.zeros((100, 100, 3), dtype=np.uint8)
-        color_display[:, :] = color_rgb[::-1]  # RGB to BGR for OpenCV
-        st.image(color_display, caption="Selected Color", width=200)
+        # Cache color display to prevent regeneration
+        if (st.session_state.color_display_cache is None or 
+            st.session_state.get('cached_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))
+            color_display = np.zeros((100, 100, 3), dtype=np.uint8)
+            color_display[:, :] = color_rgb[::-1]  # RGB to BGR for OpenCV
+            st.session_state.color_display_cache = color_display
+            st.session_state.cached_color = st.session_state.bg_color
+        
+        st.image(st.session_state.color_display_cache, caption="Selected Color", width=200)
     
     return bg_type.lower()
 
@@ -194,9 +216,11 @@ def process_video(input_file, background, bg_type="image"):
                 logger.info(f"Background color image written to {bg_path}")
 
             logger.info(f"Disk free before processing: {os.popen('df -h / | tail -1').read().strip()}")
+            
             # 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_placeholder.progress(progress)
@@ -268,30 +292,34 @@ def main():
     with col2:
         st.header("2. Background Settings")
         bg_type = handle_background_selection()
+        
         st.header("3. Process & Download")
-        # ------- FORM WRAP START --------
+        # Form wrap to prevent reruns during processing
         with st.form("process_form"):
             submitted = st.form_submit_button(
                 "🚀 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
                 with st.spinner("Processing video (this may take a few minutes)..."):
                     try:
                         # Prepare background based on type
                         background = None
-                        if bg_type == "image" and 'bg_image' in st.session_state and st.session_state.bg_image is not None:
-                            background = st.session_state.bg_image
+                        if bg_type == "image" and st.session_state.bg_image_cache is not None:
+                            background = st.session_state.bg_image_cache
                         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!")
@@ -302,9 +330,9 @@ def main():
                         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:
+        
+        # Show processed video if available (outside form for stability)
+        if st.session_state.processed_video_path:
             st.markdown("### Processed Video")
             try:
                 with open(st.session_state.processed_video_path, 'rb') as f:
@@ -315,7 +343,8 @@ def main():
                         data=video_bytes,
                         file_name="processed_video.mp4",
                         mime="video/mp4",
-                        use_container_width=True
+                        use_container_width=True,
+                        key="download_button"
                     )
             except Exception as e:
                 st.error(f"Error displaying video: {str(e)}")
@@ -323,4 +352,4 @@ def main():
 
 # --- Entry Point ---
 if __name__ == "__main__":
-    main()
+    main()