Update ui.py

ui.py

@@ -1,9 +1,8 @@
 #!/usr/bin/env python3
 """
 Modern UI for Video Background Replacer (PRO)
-- Clean, responsive Streamlit layout
-- Shows live stage progress and all errors during processing
-- Progress bars and visual feedback to keep session alive
+- File-based progress polling to keep session alive
+- Real-time progress updates during long-running operations
 """
 import streamlit as st
 import os
@@ -12,9 +11,12 @@
 import numpy as np
 import logging
 import time
+import threading
+
+from utils.progress_tracker import get_progress
 
 logger = logging.getLogger("Advanced Video Background Replacer")
-UI_BUILD = "ui-2025-10-04-16-00Z"
+UI_BUILD = "ui-2025-10-04-17-00Z"
 
 def tail_file(path: str, lines: int = 400) -> str:
     if not os.path.exists(path):
@@ -107,7 +109,6 @@ def _render_background_settings():
 
 def render_ui(process_video_func):
     try:
-        # --- Sidebar: System Status & Logs ---
         with st.sidebar:
             st.subheader("System Status")
             st.caption(f"UI build: {UI_BUILD}")
@@ -139,10 +140,8 @@ def render_ui(process_video_func):
                 log_text = tail_file("/tmp/app.log", st.session_state.get('log_tail_lines', 400))
                 st.code(log_text, language="text")
 
-        # --- Main UI: Two-Column Layout ---
         col1, col2 = st.columns([1, 1], gap="large")
 
-        # --- Column 1: Video Upload ---
         with col1:
             st.header("1. Upload Video")
             uploaded_video = st.file_uploader(
@@ -165,12 +164,10 @@ def render_ui(process_video_func):
             else:
                 video_preview_placeholder.empty()
 
-        # --- Column 2: Background & Processing ---
         with col2:
             background, bg_type = _render_background_settings()
             st.header("3. Process Video")
             
-            # Progress tracking UI elements
             progress_container = st.container()
             with progress_container:
                 progress_bar = st.progress(0)
@@ -183,59 +180,10 @@ def render_ui(process_video_func):
                 and (background is not None)
             )
 
-            # Enhanced progress callback with visual updates
-            progress_state = {"step": 0, "total_steps": 12}
-            
-            def progress_callback(msg):
-                status_text.info(f"**Status:** {msg}")
-                logger.info(f"[STAGE] {msg}")
-                
-                # Update progress based on message content
-                if "Stage 1 initiated" in msg:
-                    progress_state["step"] = 1
-                    stage_status.markdown("**Current Stage:** Stage 1 - Transparent Video Creation")
-                elif "GPU engaged" in msg or "SAM2" in msg:
-                    progress_state["step"] = 2
-                elif "SAM2 processing frame" in msg:
-                    progress_state["step"] = 3
-                elif "SAM2 complete" in msg or "clearing GPU" in msg:
-                    progress_state["step"] = 4
-                elif "MatAnyone starting" in msg:
-                    progress_state["step"] = 5
-                    stage_status.markdown("**Current Stage:** Stage 1 - MatAnyone Processing")
-                elif "MatAnyone processing" in msg:
-                    progress_state["step"] = 6
-                elif "MatAnyone complete" in msg:
-                    progress_state["step"] = 7
-                elif "Smoothing" in msg:
-                    progress_state["step"] = 8
-                elif "transparent video" in msg and "Creating" in msg:
-                    progress_state["step"] = 9
-                elif "Stage 1 complete" in msg:
-                    progress_state["step"] = 10
-                elif "Stage 2 begun" in msg:
-                    progress_state["step"] = 10
-                    stage_status.markdown("**Current Stage:** Stage 2 - Background Compositing")
-                elif "Compositing frames" in msg:
-                    progress_state["step"] = 11
-                elif "Stage 2 complete" in msg:
-                    progress_state["step"] = 12
-                    stage_status.markdown("**Current Stage:** Complete!")
-                
-                # Update progress bar
-                progress_pct = min(100, int((progress_state["step"] / progress_state["total_steps"]) * 100))
-                progress_bar.progress(progress_pct)
-
             if st.button("Process Video", disabled=not can_process, use_container_width=True):
                 try:
                     logger.info("Process Video button clicked")
                     
-                    # Reset progress
-                    progress_state["step"] = 0
-                    progress_bar.progress(0)
-                    status_text.empty()
-                    stage_status.empty()
-
                     import io
                     class _MemFile:
                         def __init__(self, name, data):
@@ -250,24 +198,58 @@ def seek(self, pos):
                     st.session_state.processing = True
                     mem_video = _MemFile("uploaded.mp4", st.session_state.video_bytes_cache)
 
-                    with st.spinner("Processing video..."):
-                        success = process_video_func(mem_video, background, bg_type.lower(), progress_callback)
+                    # Start processing in background thread
+                    thread = threading.Thread(
+                        target=process_video_func,
+                        args=(mem_video, background, bg_type.lower()),
+                        daemon=True
+                    )
+                    thread.start()
 
-                    st.session_state.processing = False
-                    if success:
+                    # Poll for progress updates
+                    while thread.is_alive() or st.session_state.get('processing', False):
+                        status = get_progress()
+                        
+                        if status.get('active') or not status.get('complete'):
+                            # Update UI
+                            progress_bar.progress(status.get('progress', 0))
+                            status_text.info(f"**Status:** {status.get('message', 'Processing...')}")
+                            stage_status.markdown(f"**Current Stage:** {status.get('stage', 'Unknown')}")
+                            
+                            # Check for errors
+                            if status.get('error'):
+                                status_text.error(f"**Error:** {status['error']}")
+                                break
+                                
+                            # Check if complete
+                            if status.get('complete'):
+                                break
+                                
+                        time.sleep(1)  # Poll every second
+                        
+                    # Wait for thread to finish
+                    thread.join(timeout=5)
+                    
+                    # Final status check
+                    final_status = get_progress()
+                    if final_status.get('error'):
+                        progress_bar.progress(final_status.get('progress', 0))
+                        status_text.error(f"**Status:** Processing failed - {final_status['error']}")
+                        st.error(f"Processing failed: {final_status['error']}")
+                    elif st.session_state.get('processed_video_bytes'):
                         progress_bar.progress(100)
                         status_text.success("**Status:** Processing complete!")
                         st.success("Video processing complete!")
                     else:
-                        status_text.error("**Status:** Processing failed. Check logs for details.")
-                        st.error("Video processing failed. Check logs for details.")
+                        status_text.error("**Status:** Processing failed. Check logs.")
+                        st.error("Processing failed. Check logs for details.")
+                        
                 except Exception as e:
                     st.session_state.processing = False
                     logger.error(f"[UI] Process video error: {e}", exc_info=True)
                     status_text.error(f"**Status:** Error - {str(e)}")
                     st.error(f"Processing error: {str(e)}. Check logs for details.")
 
-        # --- Results Display ---
         if st.session_state.get('processed_video_bytes') is not None:
             st.markdown("---")
             st.markdown("### Processed Video")
@@ -284,7 +266,6 @@ def seek(self, pos):
                 logger.error(f"[UI] Display error: {e}", exc_info=True)
                 st.error(f"Display error: {e}")
 
-        # --- Error Display ---
         if st.session_state.get('last_error'):
             with st.expander("Last Error", expanded=True):
                 st.error(st.session_state.last_error)